diff --git a/CMakeLists.txt b/CMakeLists.txt index 48f53249b0..6382018696 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -755,6 +755,57 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" AND SM70_TURBOMIND_ARCHS) "${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so") endif() +# Native SM70 FlashInfer adapters live in a separate, opt-in fragment. This +# does not relink the existing model/quantization kernels or require a user JIT. +if(VLLM_GPU_LANG STREQUAL "CUDA" AND SM70_TURBOMIND_ARCHS) + set(VLLM_SM70_FLASHINFER_SRCS "csrc/flashinfer_sm70/qsa_mqa.cu") + set_gencode_flags_for_srcs( + SRCS "${VLLM_SM70_FLASHINFER_SRCS}" + CUDA_ARCHS "${SM70_TURBOMIND_ARCHS}") + define_extension_target( + _sm70_flashinfer_C + DESTINATION vllm + LANGUAGE CUDA + SOURCES ${VLLM_SM70_FLASHINFER_SRCS} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES "${CMAKE_CURRENT_SOURCE_DIR}/flashinfer-sm70/include" + USE_SABI 3 + WITH_SOABI) + set_target_properties(_sm70_flashinfer_C PROPERTIES + CXX_STANDARD 17 CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS OFF + CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF) + target_compile_definitions(_sm70_flashinfer_C PRIVATE TORCH_API_INCLUDE_EXTENSION_H) + target_link_libraries(_sm70_flashinfer_C PRIVATE + "${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so") + + # Independent fragment: component overrides cannot double-register an + # existing operator when the MQA-only library is already loaded. + set(VLLM_SM70_FLASHINFER_GDN_SRCS + "csrc/flashinfer_sm70/gdn_h2560_q4_v12.cu" + "csrc/flashinfer_sm70/gdn_h2560_q8_v24.cu" + "csrc/flashinfer_sm70/gdn_h2560_q16_v48.cu") + set_gencode_flags_for_srcs( + SRCS "${VLLM_SM70_FLASHINFER_GDN_SRCS}" + CUDA_ARCHS "${SM70_TURBOMIND_ARCHS}") + define_extension_target( + _sm70_flashinfer_gdn_C + DESTINATION vllm + LANGUAGE CUDA + SOURCES ${VLLM_SM70_FLASHINFER_GDN_SRCS} + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES ${VLLM_GPU_ARCHES} + INCLUDE_DIRECTORIES "${CMAKE_CURRENT_SOURCE_DIR}/flashinfer-sm70/include" + USE_SABI 3 + WITH_SOABI) + set_target_properties(_sm70_flashinfer_gdn_C PROPERTIES + CXX_STANDARD 17 CXX_STANDARD_REQUIRED ON CXX_EXTENSIONS OFF + CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF) + target_compile_definitions(_sm70_flashinfer_gdn_C PRIVATE TORCH_API_INCLUDE_EXTENSION_H) + target_link_libraries(_sm70_flashinfer_gdn_C PRIVATE + "${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so") +endif() + if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") # # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY) diff --git a/MANIFEST.in b/MANIFEST.in index f9195e1c99..d875eb9823 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -7,4 +7,5 @@ include CMakeLists.txt recursive-include cmake * recursive-include csrc * +recursive-include flashinfer-sm70/include *.cuh recursive-include flash_qla *.py *.cu LICENSE diff --git a/benchmarks/benchmark_sm70_batch_tool_quality.py b/benchmarks/benchmark_sm70_batch_tool_quality.py new file mode 100644 index 0000000000..a40197bd9e --- /dev/null +++ b/benchmarks/benchmark_sm70_batch_tool_quality.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Concurrent BFCL/JSONSchemaBench API gate using the existing task scorers. + +Run identical cases/seeds against control and candidate servers separately. +Client concurrency is NOT proof of a GPU batch width or a throughput metric: +confirm actual batch dispatch from worker logs/traces. No tools are executed. +""" + +import argparse +import hashlib +import json +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import jsonschema + +from benchmarks.benchmark_sm70_tool_protocol import ( + _adapt_openai_json_schema, + _bfcl_tools, + _read_jsonl, + _stream_chat, + _validate_bfcl_calls, +) + + +def load_cases(bfcl_dir, schema_dir, per_category, schema_limit, max_schema_bytes): + cases = [] + sources = {} + for category in ("simple_python", "parallel", "multiple", "irrelevance"): + path = bfcl_dir / f"BFCL_v4_{category}.json" + sources[str(path)] = hashlib.sha256(path.read_bytes()).hexdigest() + truth = {} + if category != "irrelevance": + answers = bfcl_dir / "possible_answer" / path.name + sources[str(answers)] = hashlib.sha256(answers.read_bytes()).hexdigest() + truth = {row["id"]: row["ground_truth"] for row in _read_jsonl(answers)} + for entry in _read_jsonl(path)[:per_category]: + cases.append( + { + "id": entry["id"], + "suite": f"bfcl/{category}", + "entry": entry, + "ground_truth": ( + None if category == "irrelevance" else truth[entry["id"]] + ), + "irrelevance": category == "irrelevance", + "request": { + "messages": entry["question"][0], + "tools": _bfcl_tools(entry["function"]), + "tool_choice": "auto", + "parallel_tool_calls": True, + }, + } + ) + paths = sorted( + (p for p in schema_dir.glob("*.json") if p.stat().st_size <= max_schema_bytes), + key=lambda p: (p.stat().st_size, p.name), + ) + if len(paths) > schema_limit: + paths = ( + [paths[len(paths) // 2]] + if schema_limit == 1 + else [ + paths[round(i * (len(paths) - 1) / (schema_limit - 1))] + for i in range(schema_limit) + ] + ) + if not paths: + raise ValueError("No JSONSchemaBench cases selected") + for path in paths: + sources[str(path)] = hashlib.sha256(path.read_bytes()).hexdigest() + schema = _adapt_openai_json_schema(json.loads(path.read_text())) + jsonschema.validators.validator_for(schema).check_schema(schema) + cases.append( + { + "id": path.stem, + "suite": "json_schema", + "schema": schema, + "request": { + "messages": [ + { + "role": "system", + "content": "Generate a JSON object matching the schema.", + }, + {"role": "user", "content": json.dumps(schema)}, + ], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": path.stem, + "strict": True, + "schema": schema, + }, + }, + }, + } + ) + if len({case["id"] for case in cases}) != len(cases): + raise ValueError("Duplicate case IDs") + return cases, sources + + +def score_case(case, response): + errors = [] + if not response.get("ok"): + return ["request failed"] + if response.get("finish_reason") not in ("stop", "tool_calls"): + errors.append(f"incomplete output: {response.get('finish_reason')!r}") + if response.get("tool_calls") and response.get("finish_reason") != "tool_calls": + errors.append("tool calls have an incorrect finish_reason") + if case["suite"] == "json_schema": + try: + jsonschema.validate(json.loads(response["content"]), case["schema"]) + except (KeyError, json.JSONDecodeError, jsonschema.ValidationError) as exc: + errors.append(str(exc)) + else: + errors.extend( + _validate_bfcl_calls( + response, + case["entry"], + case["ground_truth"], + irrelevance=case["irrelevance"], + ) + ) + return errors + + +def run_cases(cases, base_url, common, concurrency, request=_stream_chat): + lock = threading.Lock() + active = peak = 0 + started = time.perf_counter() + + def run(item): + nonlocal active, peak + index, case = item + payload = {**common, **case["request"], "seed": common["seed"] + index} + with lock: + active += 1 + peak = max(peak, active) + start = time.perf_counter() - started + try: + response = request(base_url, payload) + except Exception as exc: + # Retain transport/parser failures as failed cases, never silently + # drop them or retry into a different quality sample. + response = {"ok": False, "error": f"{type(exc).__name__}: {exc}"} + finally: + end = time.perf_counter() - started + with lock: + active -= 1 + return { + "id": case["id"], + "suite": case["suite"], + "payload": payload, + "client_start_seconds": start, + "client_end_seconds": end, + "response": response, + "errors": score_case(case, response), + } + + with ThreadPoolExecutor(max_workers=concurrency) as pool: + results = list(pool.map(run, enumerate(cases))) + suites = {} + for row in results: + counts = suites.setdefault(row["suite"], {"correct": 0, "total": 0}) + counts["total"] += 1 + counts["correct"] += not row["errors"] + return { + "requested_client_concurrency": concurrency, + "peak_inflight_client_requests": peak, + "elapsed_seconds": time.perf_counter() - started, + "note": "Fixed dataset subset, not official leaderboard or speed results.", + "suites": suites, + "cases": results, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--bfcl-dir", type=Path, required=True) + parser.add_argument("--schema-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--concurrency", type=int, default=16) + parser.add_argument("--per-category", type=int, default=16) + parser.add_argument("--schema-limit", type=int, default=16) + parser.add_argument("--max-schema-bytes", type=int, default=4096) + parser.add_argument("--max-tokens", type=int, default=16384) + parser.add_argument("--seed", type=int, default=20260905) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=20) + parser.add_argument("--top-p", type=float, default=0.95) + parser.add_argument("--enable-thinking", action="store_true") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + if min(args.concurrency, args.per_category, args.schema_limit, args.max_tokens) < 1: + parser.error("Concurrency, counts and max-tokens must be positive") + cases, sources = load_cases( + args.bfcl_dir, + args.schema_dir, + args.per_category, + args.schema_limit, + args.max_schema_bytes, + ) + common = { + "model": args.model, + "stream": True, + "return_token_ids": True, + "temperature": args.temperature, + "top_p": args.top_p, + "top_k": args.top_k, + "seed": args.seed, + "max_tokens": args.max_tokens, + "chat_template_kwargs": {"enable_thinking": args.enable_thinking}, + } + result = ( + {"dry_run": True, "selected_cases": cases} + if args.dry_run + else run_cases(cases, args.base_url, common, args.concurrency) + ) + result.update(sources=sources, common_payload=common) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result.get("suites", {"selected": len(cases)}))) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_sm70_tool_protocol.py b/benchmarks/benchmark_sm70_tool_protocol.py index 00aba5f818..706d665dfe 100644 --- a/benchmarks/benchmark_sm70_tool_protocol.py +++ b/benchmarks/benchmark_sm70_tool_protocol.py @@ -732,6 +732,50 @@ def _bfcl_normalized_value(value: Any) -> Any: return value +def _bfcl_dict_matches(actual: dict[str, Any], expected: dict[str, Any]) -> bool: + # BFCL's dict_checker encodes each dictionary value as a list of acceptable + # alternatives, NOT as the literal value. Empty string permits an omitted + # key. Preserve literal lists within an alternative (do not flatten them). + return all( + key in expected + and any( + _bfcl_normalized_value(value) == _bfcl_normalized_value(candidate) + for candidate in expected[key] + ) + for key, value in actual.items() + ) and all(key in actual or "" in allowed for key, allowed in expected.items()) + + +def _bfcl_argument_matches(value: Any, allowed: list[Any], schema: dict) -> bool: + # Mirror BFCL's dict/list-of-dicts value rules; other arguments retain the + # existing comparison. This helper is not the full official AST evaluator. + if schema.get("type") in ("dict", "object") and isinstance(value, dict): + return any( + isinstance(candidate, dict) and _bfcl_dict_matches(value, candidate) + for candidate in allowed + ) + if ( + schema.get("type") == "array" + and schema.get("items", {}).get("type") in ("dict", "object") + and isinstance(value, list) + ): + return any( + isinstance(candidate, list) + and len(value) == len(candidate) + and all( + isinstance(item, dict) + and isinstance(answer, dict) + and _bfcl_dict_matches(item, answer) + for item, answer in zip(value, candidate) + ) + for candidate in allowed + ) + return any( + _bfcl_normalized_value(value) == _bfcl_normalized_value(candidate) + for candidate in allowed + ) + + def _bfcl_call_matches( actual: dict[str, Any], expected: dict[str, Any], @@ -755,10 +799,7 @@ def _bfcl_call_matches( if name not in properties or name not in expected_arguments: return False, f"unexpected argument {name!r}" allowed = expected_arguments[name] - if not any( - _bfcl_normalized_value(value) == _bfcl_normalized_value(candidate) - for candidate in allowed - ): + if not _bfcl_argument_matches(value, allowed, properties[name]): return False, f"argument {name!r} value {value!r} not in {allowed!r}" for name, allowed in expected_arguments.items(): if name not in arguments and "" not in allowed: diff --git a/benchmarks/csrc/sm70_flashinfer_gdn_conv.cu b/benchmarks/csrc/sm70_flashinfer_gdn_conv.cu new file mode 100644 index 0000000000..2be7138cb7 --- /dev/null +++ b/benchmarks/csrc/sm70_flashinfer_gdn_conv.cu @@ -0,0 +1,4 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Geometry/experimental defines belong to the benchmark build invocation. +#include "../../csrc/flashinfer_sm70/gdn_bridge.cuh" diff --git a/benchmarks/csrc/sm70_flashinfer_hc_norm.cu b/benchmarks/csrc/sm70_flashinfer_hc_norm.cu new file mode 100644 index 0000000000..0fb93b9812 --- /dev/null +++ b/benchmarks/csrc/sm70_flashinfer_hc_norm.cu @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include +#include +#include +#include +#include + +namespace { +template +void launch(torch::Tensor residual, torch::Tensor block, + torch::Tensor injection, torch::Tensor weight, + torch::Tensor combined, torch::Tensor output, float eps, int warps, + bool registers) { + const int groups = injection.size(1), d = block.size(1); + if (registers) { +#define REG_LAUNCH(NW) \ + flashinfer::sm70::hc::HCCombineNormRegisterKernel<8, 2560, NW, T, B, W> \ + <<>>( \ + (const B*)block.data_ptr(), (const T*)residual.data_ptr(), \ + (const W*)weight.data_ptr(), (const W*)injection.data_ptr(), \ + (T*)combined.data_ptr(), (T*)output.data_ptr(), groups, \ + weight.numel() == d, eps) + if (warps == 4) { + REG_LAUNCH(4); + } else { + REG_LAUNCH(8); + } +#undef REG_LAUNCH + return; + } + const int rounds = (d + 8 * 32 * warps - 1) / (8 * 32 * warps); + const int shared = + ((warps + 3) / 4 * 4 + rounds * 8 * 32 * warps) * sizeof(float); + flashinfer::sm70::hc::HCCombineNormKernel<8, T, B, W> + <<>>( + (const B*)block.data_ptr(), (const T*)residual.data_ptr(), + (const W*)weight.data_ptr(), (const W*)injection.data_ptr(), + (T*)combined.data_ptr(), (T*)output.data_ptr(), groups, d, + block.stride(0), residual.stride(0), injection.stride(0), + weight.numel() == d, eps); +} + +template +void weight_dispatch(torch::Tensor r, torch::Tensor b, torch::Tensor i, + torch::Tensor w, torch::Tensor c, torch::Tensor o, + float eps, int warps, bool registers) { + if (w.scalar_type() == at::kHalf) + launch(r, b, i, w, c, o, eps, warps, registers); + else + launch(r, b, i, w, c, o, eps, warps, registers); +} + +void run(torch::Tensor r, torch::Tensor b, torch::Tensor i, torch::Tensor w, + torch::Tensor c, torch::Tensor o, double eps, int64_t warps, + bool registers) { + const c10::cuda::CUDAGuard guard(r.device()); + for (const auto& t : {r, b, i, w, c, o}) { + TORCH_CHECK(t.is_cuda() && t.device() == r.device() && t.is_contiguous()); + TORCH_CHECK(t.scalar_type() == at::kHalf || t.scalar_type() == at::kFloat); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0); + } + const auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0); + TORCH_CHECK(r.dim() == 2 && b.dim() == 2 && i.dim() == 2 && w.dim() == 1); + TORCH_CHECK(r.size(0) == b.size(0) && r.size(0) == i.size(0)); + TORCH_CHECK(i.size(1) > 0 && b.size(1) > 0 && b.size(1) <= 4096 && + b.size(1) % 8 == 0); + TORCH_CHECK(r.size(1) == i.size(1) * b.size(1)); + TORCH_CHECK(w.numel() == b.size(1) || w.numel() == r.size(1)); + TORCH_CHECK(i.scalar_type() == w.scalar_type()); + TORCH_CHECK(c.sizes() == r.sizes() && o.sizes() == r.sizes() && + c.scalar_type() == r.scalar_type() && + o.scalar_type() == r.scalar_type()); + TORCH_CHECK(warps == 1 || warps == 2 || warps == 4 || warps == 8); + TORCH_CHECK(!registers || (b.size(1) == 2560 && (warps == 4 || warps == 8))); + TORCH_CHECK(eps > 0 && r.size(0) > 0); + if (r.scalar_type() == at::kHalf) { + if (b.scalar_type() == at::kHalf) + weight_dispatch(r, b, i, w, c, o, eps, warps, registers); + else + weight_dispatch(r, b, i, w, c, o, eps, warps, registers); + } else { + if (b.scalar_type() == at::kHalf) + weight_dispatch(r, b, i, w, c, o, eps, warps, registers); + else + weight_dispatch(r, b, i, w, c, o, eps, warps, registers); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace +TORCH_LIBRARY_FRAGMENT(_C_flashinfer_hc_sm70, m) { + m.def( + "run(Tensor residual, Tensor block, Tensor injection, Tensor weight, " + "Tensor(a!) combined, Tensor(b!) output, float eps, int warps, bool " + "registers) -> ()"); +} +TORCH_LIBRARY_IMPL(_C_flashinfer_hc_sm70, CUDA, m) { m.impl("run", &run); } diff --git a/benchmarks/csrc/sm70_flashinfer_qsa_decode.cu b/benchmarks/csrc/sm70_flashinfer_qsa_decode.cu new file mode 100644 index 0000000000..1100efc529 --- /dev/null +++ b/benchmarks/csrc/sm70_flashinfer_qsa_decode.cu @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include +#include +#include +#include + +#if defined(FI_QSA_WMMA_COMPAT) && FI_QSA_WMMA_COMPAT + #include + #define FI_QSA_NAMESPACE _C_flashinfer_qsa_sm70_compat +#else + #include + #define FI_QSA_NAMESPACE _C_flashinfer_qsa_sm70 +#endif + +namespace { +using namespace flashinfer::attention::sm70; + +void run(torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor indices, torch::Tensor table, torch::Tensor requests, + torch::Tensor offsets, torch::Tensor metadata, torch::Tensor zero, + torch::Tensor partial, torch::Tensor lse, torch::Tensor output, + int64_t splits) { + const c10::cuda::CUDAGuard guard(q.device()); + for (const auto& t : {q, k, v, indices, table, requests, offsets, metadata, + zero, partial, lse, output}) { + TORCH_CHECK(t.is_cuda() && t.device() == q.device()); + } + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major == 7 && + at::cuda::getCurrentDeviceProperties()->minor == 0, + "This benchmark entry is SM70 only; upstream gates unchanged"); + TORCH_CHECK(q.dim() == 3 && k.dim() == 4 && v.sizes() == k.sizes()); + TORCH_CHECK(q.scalar_type() == at::kHalf && k.scalar_type() == at::kHalf && + v.scalar_type() == at::kHalf && + output.scalar_type() == at::kHalf && + zero.scalar_type() == at::kHalf); + TORCH_CHECK(q.size(2) == 256 && k.size(3) == 256 && q.stride(2) == 1 && + k.stride(3) == 1 && v.strides() == k.strides()); + for (const auto& t : {q, k, v}) { + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, + "Vector loads require 16-byte aligned inputs"); + for (int axis = 0; axis + 1 < t.dim(); ++axis) + TORCH_CHECK(t.stride(axis) % 8 == 0, + "Vector loads require aligned row/head strides"); + } + TORCH_CHECK(q.stride(0) <= UINT32_MAX && q.stride(1) <= UINT32_MAX); + const int rows = q.size(0), heads = q.size(1), kv_heads = k.size(2); + TORCH_CHECK(rows > 0 && heads > 0 && heads <= 32 && kv_heads > 0 && + heads % kv_heads == 0 && k.size(0) > 0 && k.size(1) > 0); + const int group = heads / kv_heads; + TORCH_CHECK(group == 1 || group == 2 || group == 4 || group == 6 || + group == 8); + TORCH_CHECK(indices.dim() == 2 && indices.size(0) == rows && + indices.size(1) > 0 && indices.stride(1) == 1); + TORCH_CHECK(table.dim() == 2 && table.size(0) > 0 && table.size(1) > 0 && + table.stride(1) == 1); + TORCH_CHECK(requests.dim() == 1 && requests.numel() == rows && + requests.is_contiguous()); + for (const auto& t : {indices, table, requests, metadata}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(splits > 0 && splits <= 64); + const int selected = indices.size(1); + const int width = ((selected + splits - 1) / splits) * splits; + TORCH_CHECK(int64_t(rows) * width < INT32_MAX); + TORCH_CHECK(offsets.scalar_type() == at::kLong && offsets.is_contiguous() && + offsets.numel() == int64_t(rows) * width); + TORCH_CHECK(metadata.is_contiguous() && + metadata.numel() == rows + 2 + 2 * rows * splits); + TORCH_CHECK(zero.is_contiguous() && zero.numel() == 256); + TORCH_CHECK(partial.scalar_type() == at::kFloat && partial.is_contiguous() && + partial.numel() == int64_t(rows) * splits * heads * 256); + TORCH_CHECK(lse.scalar_type() == at::kFloat && lse.is_contiguous() && + lse.numel() == int64_t(rows) * splits * heads); + TORCH_CHECK(output.sizes() == q.sizes() && output.is_contiguous()); + const auto stream = at::cuda::getCurrentCUDAStream(q.get_device()); + auto* meta = metadata.data_ptr(); + PrepareQSA<<>>( + indices.data_ptr(), table.data_ptr(), + requests.data_ptr(), offsets.data_ptr(), meta, rows, + selected, width, splits, k.size(1), table.size(1), table.size(0), + k.size(0), indices.stride(0), table.stride(0), k.stride(0), k.stride(1)); + QSAParams p{}; + p.q = reinterpret_cast(q.data_ptr()); + p.o = partial.data_ptr(); + p.lse = lse.data_ptr(); + p.paged_kv.batch_size = rows; + p.paged_kv.num_heads = kv_heads; + p.paged_kv.width = width; + p.paged_kv.head_stride = k.stride(2); + p.paged_kv.k_data = {reinterpret_cast(k.data_ptr()), + reinterpret_cast(zero.data_ptr())}; + p.paged_kv.v_data = {reinterpret_cast(v.data_ptr()), + reinterpret_cast(zero.data_ptr())}; + p.paged_kv.indptr = meta; + p.paged_kv.offsets = offsets.data_ptr(); + p.padded_batch_size = rows * splits; + p.num_qo_heads = heads; + p.request_indices = meta + rows + 1; + p.kv_tile_indices = p.request_indices + rows * splits; + p.kv_chunk_size_ptr = p.kv_tile_indices + rows * splits; + p.q_stride_n = q.stride(0); + p.q_stride_h = q.stride(1); +#if defined(FI_QSA_WMMA_COMPAT) && FI_QSA_WMMA_COMPAT + #define DISPATCH(G) \ + case G: \ + LaunchQSAWMMACompatible(p, \ + reinterpret_cast(output.data_ptr()), \ + rows, splits, selected, stream); \ + break +#else + #define DISPATCH(G) \ + case G: \ + LaunchQSADecode(p, reinterpret_cast(output.data_ptr()), rows, \ + splits, stream); \ + break +#endif + switch (group) { + DISPATCH(1); + DISPATCH(2); + DISPATCH(4); + DISPATCH(6); + DISPATCH(8); + } +#undef DISPATCH + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace + +TORCH_LIBRARY_FRAGMENT(FI_QSA_NAMESPACE, m) { + m.def( + "run(Tensor q, Tensor k, Tensor v, Tensor indices, Tensor table, " + "Tensor requests, Tensor(a!) offsets, Tensor(b!) metadata, Tensor zero, " + "Tensor(c!) partial, Tensor(d!) lse, Tensor(e!) output, int splits) -> " + "()"); +} +TORCH_LIBRARY_IMPL(FI_QSA_NAMESPACE, CUDA, m) { m.impl("run", &run); } diff --git a/benchmarks/csrc/sm70_moe_compact_tasks.cu b/benchmarks/csrc/sm70_moe_compact_tasks.cu new file mode 100644 index 0000000000..93f0222443 --- /dev/null +++ b/benchmarks/csrc/sm70_moe_compact_tasks.cu @@ -0,0 +1,353 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Benchmark-only compact task scheduling; not selected by serving or CMake. +// Arithmetic and grouping copied from nvfp4_grouped_decode_sm70.cu at +// 3e0f7a40c1. Both screened schedules were negative; retain as reproducible +// evidence, without modifying the existing production grouped-MoE kernel. +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int kMaxRoutes = 160; +constexpr int kPack = 8; +constexpr int kExperts = 512; +constexpr int kChunks = kMaxRoutes / kPack; +#ifndef SM70_MOE_COMPACT_WAVES + #define SM70_MOE_COMPACT_WAVES 0 +#endif +constexpr bool kCompactTasks = SM70_MOE_COMPACT_WAVES > 0; + +// Integer atomics only: route order within a pack is immaterial because both +// projections scatter back to the original route before the unchanged W2. +__global__ void plan_kernel(const int32_t* ids, int32_t* rows, int32_t* experts, + int32_t* sizes, int32_t* total, int routes) { + __shared__ int counts[kExperts + 1]; + __shared__ int groups[kExperts + 1][kChunks]; + const int t = threadIdx.x; + for (int e = t; e <= kExperts; e += blockDim.x) counts[e] = 0; + if (t == 0) *total = 0; + __syncthreads(); + int expert = 0, ordinal = 0; + if (t < routes) { + expert = ids[t]; + if (expert < 0 || expert >= kExperts) expert = kExperts; + ordinal = atomicAdd(counts + expert, 1); + } + __syncthreads(); + if (t < routes && ordinal % kPack == 0) { + const int group = atomicAdd(total, 1); + groups[expert][ordinal / kPack] = group; + experts[group] = expert; + sizes[group] = min(kPack, counts[expert] - ordinal); + } + __syncthreads(); + if (t < routes) { + const int group = groups[expert][ordinal / kPack]; + rows[group * kPack + ordinal % kPack] = t; + } +} + +__device__ __forceinline__ void decode(unsigned packed, half2 scale, + half2* out) { + constexpr unsigned sign = 0x80008000u, em = 0x0e000e00u; + unsigned v[4] = {((packed << 12) & sign) | ((packed << 9) & em), + ((packed << 8) & sign) | ((packed << 5) & em), + ((packed << 4) & sign) | ((packed << 1) & em), + (packed & sign) | ((packed >> 3) & em)}; +#pragma unroll + for (int i = 0; i < 4; ++i) + out[i] = __hmul2(*reinterpret_cast(v + i), scale); +} + +#define PACKED_MMA(C, A0, A1, B0, B1) \ + asm volatile( \ + "mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " \ + "{%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, " \ + "{%0,%1,%2,%3,%4,%5,%6,%7};\n" \ + : "+f"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3]), "+f"(C[4]), \ + "+f"(C[5]), "+f"(C[6]), "+f"(C[7]) \ + : "r"(A0), "r"(A1), "r"(B0), "r"(B1)) + +template +__global__ void w13_kernel(const half* x, const uint32_t* weights, + const half* scales, const int32_t* rows, + const int32_t* experts, const int32_t* sizes, + const int32_t* total, half* out) { + // Split within the CTA: no floating-point atomics or global partial tensor. + __shared__ float partial[2][Split][kPack][32]; + __shared__ half projected[2][kPack][32]; + if constexpr (!kCompactTasks) { + if (blockIdx.y >= *total) return; + } + const int task_begin = kCompactTasks ? blockIdx.x : 0; + const int task_stride = kCompactTasks ? gridDim.x : 1; + const int task_end = kCompactTasks ? *total * 5 : 1; + for (int task = task_begin; task < task_end; task += task_stride) { + const int group_id = kCompactTasks ? task / 5 : blockIdx.y; + const int tile_x = kCompactTasks ? task % 5 : blockIdx.x; + const int count = sizes[group_id], expert = experts[group_id]; + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + const int projection = warp / Split, split = warp % Split; + const int tile = + Interleaved ? tile_x * 2 + projection : tile_x + projection * 5; + const int mma_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int quad = (lane >> 2) & 3; + const int col = quad * 8 + mma_row; + const int route = mma_row < count ? rows[group_id * kPack + mma_row] : 0; + float accum[8] = {}; + if (expert < kExperts) { + const uint32_t* w = weights + static_cast(expert) * 2560 * 40; + const half* s = scales + static_cast(expert) * 160 * 320; + const half* input = x + static_cast(route / 10) * 2560; +#pragma unroll 4 + for (int g = split * (160 / Split); g < (split + 1) * (160 / Split); + ++g) { + const size_t offset = + (static_cast(tile) * 320 + g * 2) * 32 + col; + const half scalar = __hmul(__ldg(s + (g * 10 + tile) * 32 + col), + __float2half_rn(16384.0f)); + const half2 scale = __halves2half2(scalar, scalar); + half2 decoded[8]; + decode(__ldcs(w + offset), scale, decoded); + decode(__ldcs(w + offset + 32), scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + uint4 lo = make_uint4(0, 0, 0, 0), hi = make_uint4(0, 0, 0, 0); + if (mma_row < count) { + lo = *reinterpret_cast(input + g * 16); + hi = *reinterpret_cast(input + g * 16 + 8); + } + PACKED_MMA(accum, lo.x, lo.y, b[0], b[1]); + PACKED_MMA(accum, lo.z, lo.w, b[2], b[3]); + PACKED_MMA(accum, hi.x, hi.y, b[4], b[5]); + PACKED_MMA(accum, hi.z, hi.w, b[6], b[7]); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int r = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + partial[projection][split][r][quad * 8 + c] = accum[i]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 2 * kPack * 32; idx += blockDim.x) { + const int p = idx / (kPack * 32), r = idx / 32 % kPack, c = idx % 32; + // FP16 materialization is retained before SiLU, then again before the + // multiplication. Split>1 changes FP32 association, not quantization. + float value = 0; +#pragma unroll + for (int s = 0; s < Split; ++s) value += partial[p][s][r][c]; + projected[p][r][c] = __float2half_rn(value); + } + __syncthreads(); + for (int idx = threadIdx.x; idx < count * 32; idx += blockDim.x) { + const int r = idx / 32, c = idx % 32; + const int p = Interleaved ? c / 16 : 0; + const int pc = Interleaved ? c % 16 * 2 : c; + const half gate = projected[p][r][pc]; + const half up = + Interleaved ? projected[p][r][pc + 1] : projected[1][r][c]; + const float gf = __half2float(gate); + const half activated = __float2half_rn(gf / (1.0f + expf(-gf))); + out[static_cast(rows[group_id * kPack + r]) * 160 + tile_x * 32 + + c] = __hmul(activated, up); + } + // A CTA may now consume another valid group/tile. Finish every output read + // before its peers reuse shared projection scratch. The existing numerical + // reduction tree and original-route scatter are unchanged. + if constexpr (kCompactTasks) __syncthreads(); + } +} + +void run(torch::Tensor out, torch::Tensor x, torch::Tensor w, torch::Tensor s, + torch::Tensor ids, torch::Tensor rows, torch::Tensor experts, + torch::Tensor sizes, torch::Tensor total, int64_t split, + bool interleaved) { + const c10::cuda::CUDAGuard guard(x.device()); + const int routes = x.size(0) * 10; + TORCH_CHECK(x.dim() == 2 && x.size(1) == 2560 && routes > 0 && routes <= 160); + for (const auto& t : {out, x, w, s, ids, rows, experts, sizes, total}) { + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous()); + } + TORCH_CHECK(x.scalar_type() == at::kHalf && s.scalar_type() == at::kHalf && + out.scalar_type() == at::kHalf && w.scalar_type() == at::kInt); + for (const auto& t : {ids, rows, experts, sizes, total}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(ids.numel() == routes && rows.numel() >= routes * kPack && + experts.numel() >= routes && sizes.numel() >= routes && + total.numel() == 1 && out.numel() == routes * 160 && + w.numel() == 512 * 2560 * 40 && s.numel() == 512 * 160 * 320); + const auto stream = at::cuda::getCurrentCUDAStream(x.get_device()); + plan_kernel<<<1, 256, 0, stream>>>( + ids.data_ptr(), rows.data_ptr(), + experts.data_ptr(), sizes.data_ptr(), + total.data_ptr(), routes); + const dim3 grid = + kCompactTasks + ? dim3(at::cuda::getCurrentDeviceProperties()->multiProcessorCount * + SM70_MOE_COMPACT_WAVES) + : dim3(5, routes); +#define LAUNCH(S, I) \ + w13_kernel<<>>( \ + reinterpret_cast(x.data_ptr()), \ + reinterpret_cast(w.data_ptr()), \ + reinterpret_cast(s.data_ptr()), rows.data_ptr(), \ + experts.data_ptr(), sizes.data_ptr(), \ + total.data_ptr(), reinterpret_cast(out.data_ptr())) +#define CASE(S) \ + case S: \ + if (interleaved) { \ + LAUNCH(S, true); \ + } else { \ + LAUNCH(S, false); \ + } \ + break + switch (split) { + CASE(1); + CASE(2); + CASE(4); + CASE(5); + CASE(8); + default: + TORCH_CHECK(false, "Unsupported split"); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +#undef CASE +#undef LAUNCH +} + +// All groups, including singletons, use one kernel. The earlier prototype +// launched separate repeated/singleton kernels; their overhead erased reuse. +__global__ void w2_kernel(const half* x, const uint32_t* weights, + const half* scales, const int32_t* rows, + const int32_t* experts, const int32_t* sizes, + const int32_t* total, half* out) { + if constexpr (!kCompactTasks) { + if (blockIdx.y * 4 + threadIdx.x / 32 >= *total) return; + } + // Keep the original CTA's four distinct expert groups and common N tile. + // Group-major warp tasks would also change weight/cache locality, confounding + // this scheduling-only experiment with the previously screened warp layout. + const int begin = kCompactTasks ? blockIdx.x : 0; + const int stride = kCompactTasks ? gridDim.x : 1; + const int end = kCompactTasks ? ((*total + 3) / 4) * 80 : 1; + for (int task = begin; task < end; task += stride) { + const int group = kCompactTasks ? task / 80 * 4 + threadIdx.x / 32 + : blockIdx.y * 4 + threadIdx.x / 32; + if constexpr (kCompactTasks) { + if (group >= *total) continue; + } + const int tile_x = kCompactTasks ? task % 80 : blockIdx.x; + const int count = sizes[group], expert = experts[group]; + const int lane = threadIdx.x % 32, quad = (lane >> 2) & 3; + const int r = (lane & 3) + ((lane & 16) ? 4 : 0), col = quad * 8 + r; + const int route = r < count ? rows[group * kPack + r] : 0; + float accum[8] = {}; + if (expert < kExperts) { + const uint32_t* w = weights + static_cast(expert) * 160 * 320; + const half* s = scales + static_cast(expert) * 10 * 2560; + const half* input = x + static_cast(route) * 160; +#pragma unroll + for (int g = 0; g < 10; ++g) { + const int offset = (tile_x * 20 + g * 2) * 32 + col; + const half scalar = __hmul(__ldg(s + (g * 80 + tile_x) * 32 + col), + __float2half_rn(16384.0f)); + const half2 scale = __halves2half2(scalar, scalar); + half2 decoded[8]; + decode(__ldcs(w + offset), scale, decoded); + decode(__ldcs(w + offset + 32), scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + uint4 lo = make_uint4(0, 0, 0, 0), hi = make_uint4(0, 0, 0, 0); + if (r < count) { + lo = *reinterpret_cast(input + g * 16); + hi = *reinterpret_cast(input + g * 16 + 8); + } + PACKED_MMA(accum, lo.x, lo.y, b[0], b[1]); + PACKED_MMA(accum, lo.z, lo.w, b[2], b[3]); + PACKED_MMA(accum, hi.x, hi.y, b[4], b[5]); + PACKED_MMA(accum, hi.z, hi.w, b[6], b[7]); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int row = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + if (row < count) { + const int dst = rows[group * kPack + row]; + out[static_cast(dst) * 2560 + tile_x * 32 + quad * 8 + c] = + __float2half_rn(accum[i]); + } + } + } +} + +__global__ void reduce_kernel(const half* routed, const float* weights, + half* out) { + const int token = blockIdx.y, col = blockIdx.x * 256 + threadIdx.x; + float result = 0; +#pragma unroll + for (int slot = 0; slot < 10; ++slot) + result = fmaf(__half2float(routed[(token * 10 + slot) * 2560 + col]), + weights[token * 10 + slot], result); + out[token * 2560 + col] = __float2half_rn(result); +} + +void w2(torch::Tensor out, torch::Tensor routed, torch::Tensor x, + torch::Tensor w, torch::Tensor s, torch::Tensor topk, + torch::Tensor rows, torch::Tensor experts, torch::Tensor sizes, + torch::Tensor total) { + const c10::cuda::CUDAGuard guard(x.device()); + TORCH_CHECK(x.dim() == 2 && x.size(1) == 160 && x.size(0) % 10 == 0); + const int routes = x.size(0), tokens = routes / 10; + TORCH_CHECK(tokens >= 1 && tokens <= 16); + for (const auto& t : + {out, routed, x, w, s, topk, rows, experts, sizes, total}) + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous()); + for (const auto& t : {out, routed, x, s}) + TORCH_CHECK(t.scalar_type() == at::kHalf); + for (const auto& t : {w, rows, experts, sizes, total}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(topk.scalar_type() == at::kFloat && topk.numel() == routes && + out.numel() == tokens * 2560 && routed.numel() == routes * 2560 && + rows.numel() >= routes * 8 && experts.numel() >= routes && + sizes.numel() >= routes && total.numel() == 1 && + w.numel() == 512 * 160 * 320 && s.numel() == 512 * 10 * 2560); + const auto stream = at::cuda::getCurrentCUDAStream(x.get_device()); + const dim3 grid = + kCompactTasks + ? dim3(at::cuda::getCurrentDeviceProperties()->multiProcessorCount * + SM70_MOE_COMPACT_WAVES) + : dim3(80, (routes + 3) / 4); + w2_kernel<<>>( + reinterpret_cast(x.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(s.data_ptr()), rows.data_ptr(), + experts.data_ptr(), sizes.data_ptr(), + total.data_ptr(), reinterpret_cast(routed.data_ptr())); + reduce_kernel<<>>( + reinterpret_cast(routed.data_ptr()), topk.data_ptr(), + reinterpret_cast(out.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace + +TORCH_LIBRARY_FRAGMENT(_C_moe_compact_tasks, m) { + m.def( + "nvfp4_grouped_w13_sm70_out(Tensor(a!) out, Tensor x, Tensor w, Tensor " + "s, Tensor ids, " + "Tensor(b!) rows, Tensor(c!) experts, Tensor(d!) sizes, Tensor(e!) " + "total, " + "int split, bool interleaved) -> ()"); + m.def( + "nvfp4_grouped_w2_sm70_out(Tensor(a!) out, Tensor(b!) routed, Tensor x, " + "Tensor w, Tensor s, " + "Tensor topk, Tensor rows, Tensor experts, Tensor sizes, Tensor total) " + "-> ()"); +} +TORCH_LIBRARY_IMPL(_C_moe_compact_tasks, CUDA, m) { + m.impl("nvfp4_grouped_w13_sm70_out", &run); + m.impl("nvfp4_grouped_w2_sm70_out", &w2); +} diff --git a/benchmarks/kernels/benchmark_qwen38_qsa_mtp5_flash_v100.py b/benchmarks/kernels/benchmark_qwen38_qsa_mtp5_flash_v100.py index f6059358d5..f0a80b8e67 100644 --- a/benchmarks/kernels/benchmark_qwen38_qsa_mtp5_flash_v100.py +++ b/benchmarks/kernels/benchmark_qwen38_qsa_mtp5_flash_v100.py @@ -44,9 +44,14 @@ def _measure_ms(call, *, warmups: int, repeats: int) -> float: def _logical_indices( *, rows: int, seq_len: int, overlap: float, seed: int, independent: bool ) -> torch.Tensor: - """Build 2051-token selections with controlled adjacent-row Page4 overlap.""" + """Build canonical 512-Page4 selections plus the open causal group tail.""" generator = torch.Generator().manual_seed(seed) - selectable_pages = torch.arange(seq_len // 4 - 1, dtype=torch.int64) + earliest_visible = seq_len if independent else seq_len - rows + 1 + # Every shared/unique page must be complete for every query. In particular, + # an MTP row must not select a future compressed page of its own block. + selectable_pages = torch.arange(earliest_visible // 4, dtype=torch.int64) + if selectable_pages.numel() < 512: + raise ValueError("each query must have at least 512 complete Page4 blocks") shared_count = round(512 * overlap) shared = selectable_pages[ torch.randperm(selectable_pages.numel(), generator=generator)[:shared_count] @@ -66,7 +71,10 @@ def _logical_indices( pages = torch.cat((shared, unique)) tokens = (pages[:, None] * 4 + torch.arange(4)).flatten() query_position = seq_len - 1 if independent else seq_len - rows + row - tail = torch.arange(query_position - 2, query_position + 1) + visible = query_position + 1 + tail = torch.full((3,), -1, dtype=torch.int64) + tail_count = visible % 4 + tail[:tail_count] = torch.arange(visible - tail_count, visible) selections.append(torch.cat((tokens, tail)).to(torch.int32)) return torch.stack(selections) diff --git a/benchmarks/kernels/benchmark_sm70_flashinfer_gdn_conv.py b/benchmarks/kernels/benchmark_sm70_flashinfer_gdn_conv.py new file mode 100644 index 0000000000..34710333f5 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_flashinfer_gdn_conv.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paired gate-projection/conv/recurrent component screen, no engine launch. + +Uses real checkpoint weights with synthetic changing hidden states. Both arms +include an identical input-refresh copy because production conv mutates QKV. +This is not a model-quality test or complete GDN layer timing. +""" + +import argparse +import hashlib +import json +import os +import statistics +import subprocess +from pathlib import Path + +import torch +from safetensors import safe_open + +from benchmarks.kernels.flashinfer_sm70_gdn_conv import FusedGDN, build + + +def check_exclusive(): + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not visible or "," in visible: + raise RuntimeError("Reserve exactly one GPU") + info = subprocess.check_output( + [ + "nvidia-smi", + "-i", + visible, + "--query-compute-apps=pid,process_name", + "--format=csv,noheader", + ], + text=True, + ) + for row in info.splitlines(): + pid, _, name = row.partition(",") + if pid.strip().isdigit() and int(pid) != os.getpid() and "snapd" not in name: + raise RuntimeError(f"Foreign GPU owner: {row}") + + +def load_weights(model, layer=0, rank=0, tp=4): + mapping = json.loads((model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + config = json.loads((model / "config.json").read_text())["text_config"] + if ( + tp < 1 + or not 0 <= rank < tp + or config["linear_num_key_heads"] % tp + or config["linear_num_value_heads"] % tp + or config["linear_key_head_dim"] != 128 + or config["linear_value_head_dim"] != 128 + ): + raise ValueError("Require D128 and a valid evenly sharded TP geometry") + prefix = f"model.language_model.layers.{layer}.linear_attn." + + def get(name, dtype=torch.float16): + name = prefix + name + with safe_open(model / mapping[name], framework="pt", device="cpu") as f: + return f.get_tensor(name).to(dtype) + + qfull = config["linear_num_key_heads"] * 128 + vfull = config["linear_num_value_heads"] * 128 + hq, hv = ( + config["linear_num_key_heads"] // tp, + config["linear_num_value_heads"] // tp, + ) + selection = torch.cat( + [ + torch.arange(rank * hq * 128, (rank + 1) * hq * 128), + torch.arange(qfull + rank * hq * 128, qfull + (rank + 1) * hq * 128), + torch.arange( + 2 * qfull + rank * hv * 128, 2 * qfull + (rank + 1) * hv * 128 + ), + ] + ) + assert qfull * 2 + vfull == get("conv1d.weight").shape[0] + wqkv = get("in_proj_qkv.weight")[selection].contiguous().cuda() + ba = ( + torch.cat( + [ + get("in_proj_b.weight")[rank * hv : (rank + 1) * hv], + get("in_proj_a.weight")[rank * hv : (rank + 1) * hv], + ] + ) + .contiguous() + .cuda() + ) + conv = get("conv1d.weight").reshape(-1, 4)[selection].contiguous().cuda() + A = get("A_log", torch.float32)[rank * hv : (rank + 1) * hv].contiguous().cuda() + dt = get("dt_bias")[rank * hv : (rank + 1) * hv].contiguous().cuda() + return config["hidden_size"], hq, hv, wqkv, ba, conv, A, dt + + +def capture(fn, repeats=1): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for _ in range(repeats): + fn() + return graph + + +def error(actual, reference): + a, b = actual.float(), reference.float() + if not a.numel(): + return {"max_abs": 0.0, "relative_l2": 0.0, "finite": True, "exact": True} + return { + "max_abs": (a - b).abs().max().item(), + "relative_l2": ((a - b).norm() / b.norm().clamp_min(1e-20)).item(), + "finite": bool(torch.isfinite(a).all()), + "exact": bool(torch.equal(actual, reference)), + } + + +@torch.inference_mode() +def screen(rows, weights, args): + from flash_qla.ops.gated_delta_rule.chunk.sm70 import fused_fwd as qla + from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update + + hidden, hq, hv, wqkv, ba_weight, cw, A, dt = weights + pool = rows + 3 + width = (2 * hq + hv) * 128 + x = torch.randn(rows, hidden, device="cuda", dtype=torch.float16) + raw = torch.empty(rows, width, device="cuda", dtype=torch.float16) + raw.copy_(x @ wqkv.t()) + base_in, candidate_in = torch.empty_like(raw), torch.empty_like(raw) + c0 = torch.randn(pool, 3, width, device="cuda", dtype=torch.float16).transpose(1, 2) + s0 = torch.randn(pool, hv, 128, 128, device="cuda", dtype=torch.float32) * 0.01 + cb, cc, sb, sc = c0.clone(), c0.clone(), s0.clone(), s0.clone() + indices = torch.arange(rows, device="cuda", dtype=torch.int32) + bias = torch.empty(0, device="cuda", dtype=torch.float16) + packed = ba_weight.t().contiguous() + ba = torch.empty(rows, 2 * hv, device="cuda", dtype=torch.float16) + b, a = ( + torch.empty(rows, hv, device="cuda", dtype=torch.float16), + torch.empty(rows, hv, device="cuda", dtype=torch.float16), + ) + out = torch.empty(rows, hv, 128, device="cuda", dtype=torch.float16) + candidate = FusedGDN(rows, hq, hv, hidden=hidden, rows_per_warp=args.rows_per_warp) + + def baseline(): + base_in.copy_(raw) + torch.mm(x, ba_weight.t(), out=ba) + b.copy_(ba[:, :hv]) + a.copy_(ba[:, hv:]) + causal_conv1d_update( + base_in, + cb, + cw, + None, + "silu", + conv_state_indices=indices, + validate_data=False, + ) + return qla.gdn_decode_mixed_qkv_global_state_sm70( + base_in, a, b, A, dt, sb, indices, out + ) + + def fused(): + candidate_in.copy_(raw) + return candidate(x, packed, candidate_in, cw, bias, cc, A, dt, sc, indices) + + baseline() + candidate_graph = capture(fused) + checks = [] + for cycle in range(args.steps): + x.normal_().mul_((0.25, 1.0, 3.0)[cycle % 3]) + raw.copy_(x @ wqkv.t()) + indices.copy_(torch.randperm(pool, device="cuda")[:rows]) + if cycle % 8 == 7: + indices[-1] = -1 + cb.copy_(c0) + cc.copy_(c0) + sb.copy_(s0) + sc.copy_(s0) + baseline() + fused() + live = indices >= 0 + checks.append( + { + "cycle": cycle, + "out": error(candidate.output[live], out[live]), + "state": error(sc, sb), + "conv_state": error(cc, cb), + "conv_out": error(candidate.conv_out[live], base_in[live]), + } + ) + eager_out, eager_state, eager_conv = ( + candidate.output.clone(), + sc.clone(), + cc.clone(), + ) + cc.copy_(c0) + sc.copy_(s0) + candidate.output.fill_(float("nan")) + candidate.partial.fill_(float("nan")) + candidate_graph.replay() + torch.cuda.synchronize() + for actual, ref in ( + (candidate.output, eager_out), + (sc, eager_state), + (cc, eager_conv), + ): + torch.testing.assert_close(actual, ref, atol=0, rtol=0) + # Advance a real recurrent history rather than always testing zero state. + c0.copy_(cb) + s0.copy_(sb) + maxima = { + part: { + metric: max(c[part][metric] for c in checks) + for metric in ("max_abs", "relative_l2") + } + for part in ("out", "state", "conv_out") + } + gate = ( + all(c[part]["finite"] for c in checks for part in ("out", "state", "conv_out")) + and all(c["conv_state"]["exact"] for c in checks) + and maxima["out"]["relative_l2"] < 5e-3 + and maxima["state"]["relative_l2"] < 5e-3 + ) + # Unlike the local comparisons above, neither arm is reset from the other + # arm in this phase. This screens accumulation through independent FP32 + # recurrent histories, including graph replays and recycled request slots. + cb.copy_(c0) + cc.copy_(c0) + sb.copy_(s0) + sc.copy_(s0) + history_maxima = { + part: {"max_abs": 0.0, "relative_l2": 0.0} + for part in ("out", "state", "conv_out") + } + history_gate, history_first_failure = True, None + for cycle in range(args.history_steps): + x.normal_().mul_((0.25, 1.0, 3.0)[cycle % 3]) + raw.copy_(x @ wqkv.t()) + indices.copy_(torch.randperm(pool, device="cuda")[:rows]) + if cycle % 8 == 7: + indices[-1] = -1 + if cycle % 31 == 30: + # A retired slot is cleared identically before being reused. Do + # not synchronize either arm's other, independently evolved slots. + slot = (cycle // 31) % pool + cb[slot].zero_() + cc[slot].zero_() + sb[slot].zero_() + sc[slot].zero_() + baseline() + candidate.output.fill_(float("nan")) + candidate.partial.fill_(float("nan")) + candidate_graph.replay() + live = indices >= 0 + current = { + "out": error(candidate.output[live], out[live]), + "state": error(sc, sb), + "conv_out": error(candidate.conv_out[live], base_in[live]), + } + for part, metrics in history_maxima.items(): + for metric in metrics: + metrics[metric] = max(metrics[metric], current[part][metric]) + current_gate = ( + all(check["finite"] for check in current.values()) + and torch.equal(cc, cb) + and bool((candidate.output[~live] == 0).all()) + and current["out"]["relative_l2"] < 5e-3 + and current["state"]["relative_l2"] < 5e-3 + ) + if not current_gate and history_first_failure is None: + history_first_failure = {"cycle": cycle, "errors": current} + history_gate = history_gate and current_gate + gate = gate and history_gate + record = { + "rows": rows, + "q_heads": hq, + "v_heads": hv, + "rows_per_warp": args.rows_per_warp, + "checks": checks, + "maxima": maxima, + "independent_history": { + "steps": args.history_steps, + "maxima": history_maxima, + "gate": history_gate, + "first_failure": history_first_failure, + }, + "operator_gate": gate, + } + if gate: + indices.copy_(torch.arange(rows, device="cuda")) + graphs = [capture(baseline, args.calls), capture(fused, args.calls)] + for _ in range(20): + for graph in graphs: + graph.replay() + torch.cuda.synchronize() + samples = [[], []] + check_exclusive() + for repeat in range(args.samples): + for i in [0, 1] if repeat % 2 == 0 else [1, 0]: + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + graphs[i].replay() + end.record() + end.synchronize() + samples[i].append(start.elapsed_time(end) * 1000 / args.calls) + check_exclusive() + record.update( + samples_us=samples, median_us=[statistics.median(s) for s in samples] + ) + record["reference_qla_binary_sha256"] = hashlib.sha256( + Path(qla._load_ext().__file__).read_bytes() + ).hexdigest() + return record + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--tp", type=int, default=4) + parser.add_argument("--rank", type=int, default=0) + parser.add_argument("--layer", type=int, default=0) + parser.add_argument("--rows-per-warp", type=int, choices=(4, 8), default=8) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 4, 8, 16]) + parser.add_argument("--steps", type=int, default=32) + parser.add_argument("--history-steps", type=int, default=256) + parser.add_argument("--samples", type=int, default=9) + parser.add_argument("--calls", type=int, default=30) + args = parser.parse_args() + if min(args.steps, args.history_steps, args.samples, args.calls) <= 0: + parser.error("Step, sample and call counts must be positive") + torch.manual_seed(20260906) + check_exclusive() + weights = load_weights(args.model, args.layer, args.rank, args.tp) + build(*weights[:3], rows_per_warp=args.rows_per_warp) + for rows in args.rows: + print(json.dumps(screen(rows, weights, args)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_flashinfer_gdn_phases.py b/benchmarks/kernels/benchmark_sm70_flashinfer_gdn_phases.py new file mode 100644 index 0000000000..ef28cc0380 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_flashinfer_gdn_phases.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Same-arithmetic GDN cooperative versus stream-ordered split phases. + +Includes BA, conv, gating and recurrence. Excludes QKVZ/output projections. +The baseline library must be the frozen integration binary, not a new rebuild. +""" + +import argparse +import hashlib +import json +from pathlib import Path +from statistics import median + +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv import ( + capture, + check_exclusive, + load_weights, +) +from benchmarks.kernels.benchmark_sm70_flashinfer_mqa import measure +from benchmarks.kernels.flashinfer_sm70_gdn_conv import FusedGDN, build +from benchmarks.kernels.sm70_paired_stats import paired_latency_interval + + +def screen(rows, weights, steps, repeats, candidate="two_phase"): + hidden, hq, hv, wqkv, ba, cw, A, dt = weights + torch.manual_seed(20260906) + x = torch.randn(rows, hidden, device="cuda", dtype=torch.float16) + raw = (x @ wqkv.T).contiguous() + pool = rows + 3 + c0 = torch.randn(pool, raw.shape[1], 3, device="cuda", dtype=torch.float16) + s0 = torch.randn(pool, hv, 128, 128, device="cuda") * 0.01 + convs, states = [c0.clone() for _ in range(2)], [s0.clone() for _ in range(2)] + indices = torch.arange(rows, device="cuda", dtype=torch.int32) + packed = ba.T.contiguous() + bias = x.new_empty(0) + ops = [ + FusedGDN(rows, hq, hv, hidden=hidden, **{candidate: p}) for p in (False, True) + ] + calls = [ + lambda i=i: ops[i]( + x, packed, raw, cw, bias, convs[i], A, dt, states[i], indices + ) + for i in range(2) + ] + graphs = [capture(call) for call in calls] + for c, s in zip(convs, states): + c.copy_(c0) + s.copy_(s0) + for step in range(steps): + x.normal_().mul_((0.25, 1.0, 3.0)[step % 3]) + raw.copy_(x @ wqkv.T) + indices.copy_(torch.randperm(pool, device="cuda")[:rows]) + if step % 8 == 7: + indices[-1] = -1 + for op, graph in zip(ops, graphs): + op.output.fill_(torch.nan) + op.partial.fill_(torch.nan) + graph.replay() + torch.cuda.synchronize() + for a, b in ( + (convs[0], convs[1]), + (states[0], states[1]), + (ops[0].output, ops[1].output), + (ops[0].partial, ops[1].partial), + ): + torch.testing.assert_close(a, b, atol=0, rtol=0) + assert torch.isfinite(a).all() + # History ends on a padding case: restore all live rows before timing. + # Otherwise B1 measures only padding and B16 really measures 15 requests. + indices.copy_(torch.arange(rows, device="cuda", dtype=torch.int32)) + times = [[], []] + for cycle in range(5): + for i in [0, 1] if cycle % 2 == 0 else [1, 0]: + graphs[i].replay() + times[i].append(measure(graphs[i].replay, repeats)) + for a, b in ( + (convs[0], convs[1]), + (states[0], states[1]), + (ops[0].output, ops[1].output), + ): + torch.testing.assert_close(a, b, atol=0, rtol=0) + assert torch.isfinite(a).all() + return dict( + rows=rows, + timed_live_rows=rows, + independent_history_steps=steps, + output_state_conv_partials_exact=True, + samples_us=times, + median_us=list(map(median, times)), + reduction_pct=100 * (1 - median(times[1]) / median(times[0])), + paired_interval=paired_latency_interval(*times), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--baseline-library", type=Path, required=True) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 4, 8, 16, 32, 64]) + parser.add_argument("--steps", type=int, default=256) + parser.add_argument( + "--candidate", choices=("two_phase", "shared_parameters"), default="two_phase" + ) + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.output.exists() or min(args.steps, args.repeats) < 1: + parser.error("Require fresh output and positive steps/repeats") + if not all(1 <= rows <= 64 for rows in args.rows): + parser.error("Rows must be in 1..64") + check_exclusive() + torch.ops.load_library(str(args.baseline_library.resolve())) + weights = load_weights(args.model) + build( + hidden=weights[0], + q_heads=weights[1], + v_heads=weights[2], + **{args.candidate: True}, + ) + results = [] + for rows in args.rows: + result = screen(rows, weights, args.steps, args.repeats, args.candidate) + print(json.dumps(result), flush=True) + results.append(result) + check_exclusive() + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps( + dict( + candidate=args.candidate, + scope=( + "native GDN BA+conv+recurrent only; not full model quality or speed" + ), + baseline_sha256=hashlib.sha256( + args.baseline_library.read_bytes() + ).hexdigest(), + results=results, + ), + indent=2, + ) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_flashinfer_hc_norm.py b/benchmarks/kernels/benchmark_sm70_flashinfer_hc_norm.py new file mode 100644 index 0000000000..0d6a43e557 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_flashinfer_hc_norm.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Complete HC combine + Gemma-norm micro, not projection or model timing.""" + +import argparse +import json +import statistics + +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv import ( + capture, + check_exclusive, + error, +) +from benchmarks.kernels.flashinfer_sm70_hc_norm import HCNorm, build + + +@torch.inference_mode() +def screen(rows, dtype, args): + from vllm.models.qwen4_exp.nvidia.ops.hc import hc_combine_norm + + r = torch.randn(rows, 4 * 2560, device="cuda", dtype=dtype) + b = torch.randn(rows, 2560, device="cuda", dtype=torch.float16) + inj = torch.randn(rows, 4, device="cuda", dtype=torch.float16) + weight = torch.randn(4 * 2560, device="cuda", dtype=torch.float16) * 0.05 + calls = [lambda: hc_combine_norm(r, b, inj, weight, 1e-6, 4)] + candidates = [HCNorm(r, warps) for warps in (1, 2, 4, 8)] + candidates += [HCNorm(r, warps, registers=True) for warps in (4, 8)] + calls += [lambda candidate=c: candidate(r, b, inj, weight) for c in candidates] + graphs = [capture(fn, args.calls) for fn in calls] + checks = [] + for cycle, scale in enumerate((0.25, 1.0, 3.0, 1.0)): + r.normal_().mul_(scale) + b.normal_().mul_(scale) + inj.normal_() + expected = calls[0]() + for index, candidate in enumerate(candidates): + eager = [x.clone() for x in calls[index + 1]()] + candidate.combined.fill_(float("nan")) + candidate.normalized.fill_(float("nan")) + graphs[index + 1].replay() + torch.cuda.synchronize() + for x, y in zip((candidate.combined, candidate.normalized), eager): + torch.testing.assert_close(x, y, atol=0, rtol=0) + diffs = [error(x, y) for x, y in zip(eager, expected)] + checks.append( + { + "cycle": cycle, + "warps": candidate.warps, + "registers": candidate.registers, + "errors": diffs, + } + ) + gate = all( + d["finite"] and d["relative_l2"] < 1e-3 for c in checks for d in c["errors"] + ) + result = { + "rows": rows, + "residual_dtype": str(dtype), + "operator_gate": gate, + "checks": checks, + "variants": [{"warps": c.warps, "registers": c.registers} for c in candidates], + } + if gate: + for _ in range(20): + for graph in graphs: + graph.replay() + torch.cuda.synchronize() + samples = [[] for _ in calls] + for repeat in range(args.samples): + order = ( + range(len(calls)) if repeat % 2 == 0 else reversed(range(len(calls))) + ) + for index in order: + s, e = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + s.record() + graphs[index].replay() + e.record() + e.synchronize() + samples[index].append(s.elapsed_time(e) * 1000 / args.calls) + check_exclusive() + result.update( + median_us=[statistics.median(s) for s in samples], samples_us=samples + ) + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", nargs="+", type=int, default=[1, 4, 8, 16]) + parser.add_argument("--calls", type=int, default=100) + parser.add_argument("--samples", type=int, default=9) + args = parser.parse_args() + check_exclusive() + torch.manual_seed(20260906) + build() + for dtype in (torch.float16, torch.float32): + for rows in args.rows: + print(json.dumps(screen(rows, dtype, args)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_flashinfer_mqa.py b/benchmarks/kernels/benchmark_sm70_flashinfer_mqa.py new file mode 100644 index 0000000000..c2748dbae0 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_flashinfer_mqa.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paired native MQA plan+score versus actual paged Triton serving arithmetic. + +Frozen 256K capacity, independent request pages and live-length changes. +This measures indexer scoring only, NOT complete QSA or endpoint performance. +""" + +import argparse +import json +import math +import os +import statistics +from pathlib import Path + +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv import check_exclusive +from benchmarks.kernels.flashinfer_sm70_mqa import FlashInferMQA, build +from benchmarks.kernels.sm70_paired_stats import paired_latency_interval +from vllm import envs +from vllm.models.qwen4_exp.nvidia.ops.qsa import _qsa_mqa_paged_kernel +from vllm.triton_utils import triton + + +def make_inputs(rows, length, dim=128, heads=4, table_width=335): + device = "cuda" + page_size = 196 + torch.manual_seed(20260906) + q = torch.randn(rows, heads, dim, device=device, dtype=torch.float16) + k = torch.randn( + rows * table_width, page_size, 1, dim, device=device, dtype=torch.float16 + ) + table = torch.randperm(rows * table_width, device=device).to(torch.int32) + table = table.reshape(rows, table_width) + requests = torch.arange(rows, device=device, dtype=torch.int32) + # QSA metadata uses int64 logical_positions in the actual model runner. + positions = torch.full_like(requests, length - 1, dtype=torch.int64) + lengths = torch.full_like(requests, length) + return q, k, table, requests, positions, lengths + + +def reference(inputs, logits, visible): + q, k, table, requests, positions, lengths = inputs + rows, heads, dim = q.shape + columns = logits.shape[1] + block_n = 32 if rows == 1 else 64 + grouping = 1 if rows <= 32 else 8 + _qsa_mqa_paged_kernel[(rows, triton.cdiv(columns, block_n * grouping))]( + q, + k, + table, + requests, + positions, + lengths, + visible, + logits, + *q.stride(), + k.stride(0), + k.stride(1), + k.stride(3), + *table.stride(), + logits.stride(0), + rows, + columns, + k.shape[0], + table.shape[0], + math.sqrt(dim), + PAGE_SIZE=k.shape[1], + PAGE_TABLE_WIDTH=table.shape[1], + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_N=block_n, + BLOCK_D=max(16, triton.next_power_of_2(dim)), + TILES_PER_PROG=grouping, + STAGES=2, + MAX_N=max(16, triton.next_power_of_2(heads)), + COMPRESS_RATIO=4, + num_warps=2, + ) + + +def check_schedule(op): + visible = op.visible.cpu().tolist() + counts = [(v + 63) // 64 for v in visible] + total = sum(counts) + quotient, remainder = divmod(total, op.workers) + expected = [] + for worker in range(op.workers + 1): + offset = worker * quotient + min(worker, remainder) + row = 0 + while row < len(counts) and offset >= counts[row]: + offset -= counts[row] + row += 1 + expected.append([row, offset]) + assert op.schedule.cpu().tolist() == expected + + +def check_scores(inputs, op, logits, visible): + assert torch.equal(op.visible, visible) + live = torch.arange(logits.shape[1], device=logits.device)[None] < visible[:, None] + a, b = op.logits[live], logits[live] + assert torch.equal(torch.isneginf(a), torch.isneginf(b)) + finite = torch.isfinite(b) + assert torch.isfinite(a[finite]).all() + error = a[finite] - b[finite] + max_abs = error.abs().max().item() if error.numel() else 0.0 + relative = (error.norm() / b[finite].norm().clamp_min(1e-12)).item() + torch.testing.assert_close(a, b, atol=2e-5, rtol=2e-5) + check_schedule(op) + return dict(max_abs=max_abs, relative_l2=relative) + + +def measure(call, repeats): + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(repeats): + call() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1000 / repeats + + +def run_case(rows, length, workers, repeats): + inputs = make_inputs(rows, length) + q, k, table, requests, positions, lengths = inputs + columns = table.shape[1] * k.shape[1] + ops = {str(w): FlashInferMQA(q, columns, w) for w in workers} + logits = torch.empty((rows, columns), device=q.device) + visible = torch.empty_like(requests) + calls = {"reference": lambda: reference(inputs, logits, visible)} + calls.update({name: lambda op=op: op(*inputs) for name, op in ops.items()}) + graphs = {} + for name, call in calls.items(): + for _ in range(3): + call() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + graphs[name] = graph + + errors = {name: [] for name in ops} + for iteration in range(6): + values = [max(0, length - row * 257 - iteration * 197) for row in range(rows)] + if iteration == 1: + values = [0 if row % 2 else 256 for row in range(rows)] + if iteration == 2: + values = [262144 - row * 3 for row in range(rows)] + if iteration == 3: + values = [255 + row % 3 for row in range(rows)] + if iteration == 4: + values = [0] * rows + lengths.copy_(torch.tensor(values, device=q.device, dtype=torch.int32)) + positions.copy_(lengths - 1) + q.mul_(-0.875) + logits.fill_(float("nan")) + for op in ops.values(): + op.logits.fill_(float("nan")) + op.schedule.fill_(-991) + for graph in graphs.values(): + graph.replay() + for name, op in ops.items(): + errors[name].append(check_scores(inputs, op, logits, visible)) + + lengths.fill_(length) + positions.fill_(length - 1) + times = {name: [] for name in graphs} + for cycle in range(5): + names = list(graphs) if cycle % 2 == 0 else list(reversed(graphs)) + for name in names: + graphs[name].replay() + times[name].append(measure(graphs[name].replay, repeats)) + for graph in graphs.values(): + graph.replay() + for op in ops.values(): + check_scores(inputs, op, logits, visible) + ref = statistics.median(times["reference"]) + return dict( + rows=rows, + live_length=length, + capacity_columns=columns, + scope="MQA plan+score only; not full QSA or model quality", + changed_replays=6, + timings={ + name: dict( + median_us=statistics.median(t), + samples_us=t, + reduction_pct=100 * (1 - statistics.median(t) / ref), + ) + for name, t in times.items() + }, + errors=errors, + paired_intervals={ + name: paired_latency_interval(times["reference"], t) + for name, t in times.items() + if name != "reference" + }, + ) + + +def screen_selector(rows, length, repeats): + from vllm.model_executor.layers import sm70_flashinfer_batch as fi + from vllm.models.qwen4_exp.nvidia.ops.qsa import qsa_select_paged_tokens + + inputs = make_inputs(rows, length) + q = inputs[0] + fi._MQA_SMS[q.device] = torch.cuda.get_device_properties( + q.device + ).multi_processor_count + saved = os.environ.get("VLLM_SM70_FLASHINFER_BATCH") + outputs = [ + torch.empty((rows, 2051), device=q.device, dtype=torch.int32) for _ in range(2) + ] + graphs = [] + try: + for mode in range(2): + os.environ["VLLM_SM70_FLASHINFER_BATCH"] = str(mode) + envs.disable_envs_cache() + for _ in range(3): + qsa_select_paged_tokens(*inputs, 2048, 4, outputs[mode]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + qsa_select_paged_tokens(*inputs, 2048, 4, outputs[mode]) + graphs.append(graph) + for cycle in range(8): + inputs[4].copy_(inputs[5] - 1 - cycle) + q.mul_(-0.875) + for output, graph in zip(outputs, graphs): + output.fill_(-919) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(outputs[0], outputs[1]), (rows, length, cycle) + inputs[4].copy_(inputs[5] - 1) + samples = [[], []] + for cycle in range(5): + for mode in [0, 1] if cycle % 2 == 0 else [1, 0]: + graphs[mode].replay() + samples[mode].append(measure(graphs[mode].replay, repeats)) + assert torch.equal(outputs[0], outputs[1]) + medians = list(map(statistics.median, samples)) + return dict( + rows=rows, + length=length, + scope="MQA+topk+index expansion; excludes sparse attention and model", + changed_replays_exact=8, + samples_us=samples, + median_us=medians, + reduction_pct=100 * (1 - medians[1] / medians[0]), + paired_interval=paired_latency_interval(*samples), + ) + finally: + if saved is None: + os.environ.pop("VLLM_SM70_FLASHINFER_BATCH", None) + else: + os.environ["VLLM_SM70_FLASHINFER_BATCH"] = saved + envs.disable_envs_cache() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", nargs="+", type=int, default=[4, 8, 16]) + parser.add_argument("--lengths", nargs="+", type=int, default=[8192, 65536, 262144]) + parser.add_argument("--workers", nargs="+", type=int, default=[80, 160, 320]) + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument( + "--selector", + action="store_true", + help="Include production topk/index expansion after the scorer", + ) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.output.exists(): + parser.error("Refusing to overwrite evidence") + if not all(1 <= b <= 64 for b in args.rows) or args.repeats < 1: + parser.error("Require positive repeats and rows in 1..64") + if not all(1 <= n <= 262144 for n in args.lengths): + parser.error("Lengths must fit the frozen capacity") + check_exclusive() + build() + results = [] + for rows in args.rows: + for length in args.lengths: + result = ( + screen_selector(rows, length, args.repeats) + if args.selector + else run_case(rows, length, args.workers, args.repeats) + ) + results.append(result) + print(json.dumps(result), flush=True) + check_exclusive() + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(dict(results=results), indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_flashinfer_qsa.py b/benchmarks/kernels/benchmark_sm70_flashinfer_qsa.py new file mode 100644 index 0000000000..bd37f5241b --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_flashinfer_qsa.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Actual FlashInfer SM70 QSA screening, not a serving performance claim. + +All candidates include index preparation and final merge. No GPU allocations +occur within the candidate forward. Run on an exclusively reserved GPU. +""" + +import argparse +import hashlib +import json +import os +import statistics +import subprocess +from pathlib import Path + +import torch + +from benchmarks.kernels.flashinfer_sm70_qsa import ( + UPSTREAM_SHA, + FlashInferQSA, + build, +) + + +def make_case(rows, selected=2051, page=784, kv_heads=1, group=6, length=8192): + """Independent requests, permuted physical pages, ordered sparse indices.""" + blocks = (length + page - 1) // page + q = torch.randn(rows, kv_heads * group, 256, device="cuda", dtype=torch.float16) + k = torch.randn( + rows * blocks, page, kv_heads, 256, device="cuda", dtype=torch.float16 + ) + v = torch.randn_like(k) + table = torch.randperm(rows * blocks, device="cuda", dtype=torch.int32) + table = table.reshape(rows, blocks) + requests = torch.arange(rows, device="cuda", dtype=torch.int32) + # Random order is deliberate, with no physical sorting or deduplication. + indices = torch.stack( + [ + torch.randperm(length, device="cuda", dtype=torch.int32)[:selected] + for _ in range(rows) + ] + ) + if selected == 2051: + # Runtime uses 512 four-token blocks and up to three tail positions. + pages = torch.stack( + [ + torch.randperm(length // 4 - 1, device="cuda", dtype=torch.int32)[:512] + for _ in range(rows) + ] + ) + indices[:, :2048] = ( + pages[:, :, None] * 4 + torch.arange(4, device="cuda") + ).reshape(rows, 2048) + indices[:, 2048:] = -1 # 8192 has zero tail residue. + return q, k, v, indices, table, requests + + +def oracle(q, k, v, indices, table, requests): + """FP32 arithmetic with the same visible-index contract, preserving repeats.""" + out = torch.zeros(q.shape, device=q.device, dtype=torch.float32) + page = k.shape[1] + group = q.shape[1] // k.shape[2] + for row in range(q.shape[0]): + request = int(requests[row]) + if not 0 <= request < table.shape[0]: + continue + logical = indices[row].long() + valid = (logical >= 0) & (logical // page < table.shape[1]) + logical = logical[valid] + physical = table[request, logical // page].long() + valid = (physical >= 0) & (physical < k.shape[0]) + logical, physical = logical[valid], physical[valid] + if logical.numel() == 0: + continue + keys = k[physical, logical % page].float().repeat_interleave(group, 1) + values = v[physical, logical % page].float().repeat_interleave(group, 1) + scores = torch.einsum("hd,shd->hs", q[row].float(), keys) / 16 + out[row] = torch.einsum("hs,shd->hd", scores.softmax(-1), values) + return out + + +def check_exclusive(): + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not visible or "," in visible: + raise RuntimeError("Choose exactly one reserved CUDA_VISIBLE_DEVICES GPU") + report = subprocess.check_output( + [ + "nvidia-smi", + "-i", + visible, + "--query-compute-apps=pid,process_name", + "--format=csv,noheader", + ], + text=True, + ) + for line in report.splitlines(): + pid, _, name = line.partition(",") + if ( + pid.strip().isdigit() + and int(pid) != os.getpid() + and "snapd-desktop-integration" not in name + ): + raise RuntimeError(f"Foreign GPU process; discard timing: {line}") + + +def capture(call, repeats): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for _ in range(repeats): + call() + return graph + + +def elapsed_us(graph, calls): + start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + start.record() + graph.replay() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1000 / calls + + +@torch.inference_mode() +def screen(rows, args): + case = make_case(rows, length=args.length) + q, k, v, indices, table, req = case + names, calls = [], [] + reference_sha = None + if args.compare_triton: + from vllm.models.qwen4_exp.nvidia.ops import qsa as ops + + # No query_positions supplied: this forces actual Triton sparse QSA, + # bypassing the native page4 dispatcher. No env heuristic override. + out = torch.empty_like(q) + calls.append( + lambda: ops.qsa_sparse_paged_attention( + q, k, v, indices, table, req, out=out + ) + ) + names.append("current_triton") + reference_sha = hashlib.sha256(Path(ops.__file__).read_bytes()).hexdigest() + candidates = [] + for splits in args.splits: + candidate = FlashInferQSA(q, indices.shape[1], splits) + candidates.append(candidate) + calls.append(lambda candidate=candidate: candidate(*case)) + names.append(f"flashinfer_s{splits}") + graphs = [capture(call, args.calls) for call in calls] + checks = [] + for cycle in range(4): + q.normal_().mul_((0.25, 1.0, 3.0, 1.0)[cycle]) + indices[:, 2048:] = -1 + if cycle: + indices[:, 2048 : 2048 + cycle] = torch.arange( + args.length - cycle, args.length, device="cuda" + ) + table.copy_(torch.randperm(k.shape[0], device="cuda").reshape_as(table)) + expected = oracle(*case) + for name, call, graph in zip(names, calls, graphs): + eager = call().clone() + graph.replay() + torch.cuda.synchronize() + # Another call is not used to inspect replay output. + replay = ( + out + if name == "current_triton" + else candidates[ + args.splits.index(int(name.removeprefix("flashinfer_s"))) + ].output + ) + torch.testing.assert_close(replay, eager, atol=0, rtol=0) + torch.testing.assert_close(replay.float(), expected, atol=2e-3, rtol=1e-2) + relative = ( + (replay.float() - expected).norm() / expected.norm().clamp_min(1e-20) + ).item() + if relative > 5e-3: + raise AssertionError(f"{name}: relative L2 {relative}") + checks.append( + { + "cycle": cycle, + "route": name, + "max_abs": (replay.float() - expected).abs().max().item(), + "relative_l2": relative, + } + ) + # Restore ordinary full-context tail for timing; all arms see same data. + indices[:, 2048:] = -1 + timings = {name: [] for name in names} + check_exclusive() + # Warm GPU clocks as well as JIT/graphs before paired sampling. Do not + # change clock policy or compare against timings from another process. + for _ in range(30): + for graph in graphs: + graph.replay() + torch.cuda.synchronize() + for iteration in range(args.samples): + order = list(range(len(names))) + if iteration % 2: + order.reverse() + for index in order: + timings[names[index]].append(elapsed_us(graphs[index], args.calls)) + check_exclusive() + hardware = subprocess.check_output( + [ + "nvidia-smi", + "-i", + os.environ["CUDA_VISIBLE_DEVICES"], + "--query-gpu=name,driver_version,clocks.current.sm,clocks.current.memory,temperature.gpu", + "--format=csv,noheader", + ], + text=True, + ).strip() + return { + "rows": rows, + "length": args.length, + "selection_width": 2051, + "q_shape": list(q.shape), + "page_size": k.shape[1], + "flashinfer_sha": UPSTREAM_SHA, + "triton_source_sha256": reference_sha, + "checks": checks, + "microseconds_samples": timings, + "median_us": {name: statistics.median(t) for name, t in timings.items()}, + "hardware_after_sampling": hardware, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", nargs="+", type=int, default=[1, 4, 8, 16]) + parser.add_argument("--splits", nargs="+", type=int, default=[16, 32, 64]) + parser.add_argument("--length", type=int, default=8192) + parser.add_argument("--samples", type=int, default=9) + parser.add_argument("--calls", type=int, default=30) + parser.add_argument("--compare-triton", action="store_true") + args = parser.parse_args() + torch.manual_seed(7) + torch.backends.cuda.matmul.allow_tf32 = False + check_exclusive() + build() + for rows in args.rows: + print(json.dumps(screen(rows, args)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_gdn_projection_split.py b/benchmarks/kernels/benchmark_sm70_gdn_projection_split.py new file mode 100644 index 0000000000..00b74eda57 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_gdn_projection_split.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fuse only the four GDN projection copies, never the GEMM arithmetic. + +The current opaque M1 op falls back to two unchanged GEMMs followed by four +contiguous copies at M>1. Screen one bit-preserving copy kernel including +both GEMMs and rotating weight allocations. No model-quality claim. +""" + +import argparse +import json +import statistics +from pathlib import Path + +import torch +from safetensors import safe_open + +from benchmarks.kernels.benchmark_sm70_moe_packed_w13 import graph, latency +from vllm.models.qwen4_exp.nvidia.sm70_fp16_gemv import ( + _qwen38_gdn_projection_split_kernel as split_kernel, +) +from vllm.models.qwen4_exp.nvidia.sm70_fp16_gemv import ( + _split_gdn_projection_outputs as split, +) + + +def reference(qkvz, ba): + return ( + qkvz[:, :2560].contiguous(), + qkvz[:, 2560:].contiguous(), + ba[:, :12].contiguous(), + ba[:, 12:].contiguous(), + ) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model", type=Path, required=True) + p.add_argument("--out", type=Path, required=True) + p.add_argument("--tokens", nargs="+", type=int, default=[2, 4, 8, 16, 32, 64]) + args = p.parse_args() + if args.out.exists() or min(args.tokens) < 2: + p.error("Output must be new; this screen is for M>=2") + assert torch.cuda.get_device_capability() == (7, 0) + torch.manual_seed(20260905) + # Exercise every 16-bit payload in each of the four store branches. + bits = torch.arange(65536, device="cuda").to(torch.int16) + exhaustive = bits.view(torch.float16)[:, None].expand(-1, 2).contiguous() + exhaustive_out = tuple(torch.empty_like(bits).view(torch.float16) for _ in range(4)) + split_kernel[(65536, 1)]( + exhaustive, + exhaustive, + *exhaustive_out, + QKV=1, + Z=1, + B=1, + A=1, + BLOCK=256, + num_warps=4, + num_stages=1, + ) + assert all(torch.equal(value.view(torch.int16), bits) for value in exhaustive_out) + index = json.loads((args.model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + weights = {} + for name in ("qkv", "z", "b", "a"): + key = f"model.language_model.layers.0.linear_attn.in_proj_{name}.weight" + with safe_open(args.model / index[key], framework="pt", device="cpu") as f: + weights[name] = f.get_tensor(key).half().cuda() + q, k, v = weights["qkv"].split((2048, 2048, 6144)) + wq = torch.cat((q[:512], k[:512], v[:1536], weights["z"][:1536])).contiguous() + wb = torch.cat((weights["b"][:12], weights["a"][:12])).contiguous() + assert wq.shape == (4096, 2560) and wb.shape == (24, 2560) + # Sixteen allocations exceed L2, but contain the same actual layer's weights. + copies = [(wq.clone(), wb.clone()) for _ in range(16)] + rows = [] + for m in args.tokens: + x = torch.randn(m, 2560, device="cuda", dtype=torch.float16) * 0.1 + qkvz = torch.empty(m, 4096, device="cuda", dtype=torch.float16) + ba = torch.empty(m, 24, device="cuda", dtype=torch.float16) + saved = [None, None] + + def run(mode, linear=False, x=x, qkvz=qkvz, ba=ba, saved=saved): + if linear: + for w, wg in copies: + q = torch.nn.functional.linear(x, w) + g = torch.nn.functional.linear(x, wg) + saved[mode] = (reference if mode == 0 else split)(q, g) + else: + saved[mode] = (reference if mode == 0 else split)(qkvz, ba) + + gs = [graph(lambda mode=mode: run(mode)) for mode in (0, 1)] + # Raw payload patterns include signed zero, subnormals, infinities, NaNs. + for shift in (0, 17, 32768, 65500): + for value in (qkvz, ba): + bits = ( + (torch.arange(value.numel(), device="cuda") + shift) % 65536 + ).to(torch.int16) + value.view(torch.int16).copy_(bits.reshape_as(value)) + for mode in (0, 1): + for value in saved[mode]: + value.view(torch.int16).fill_(12345) + gs[mode].replay() + for actual, expected in zip(saved[1], saved[0], strict=True): + assert actual.is_contiguous() + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) + for scope in ("copies", "complete_projection_rotating_weights"): + graphs = ( + gs + if scope == "copies" + else [ + graph(lambda mode=mode: run(mode, True), unroll=1) + for mode in (0, 1) + ] + ) + for mode in (0, 1): + graphs[mode].replay() + assert all( + torch.equal(a.view(torch.int16), b.view(torch.int16)) + for a, b in zip(saved[0], saved[1], strict=True) + ) + times = [[], []] + for sample in range(5): + for mode in (0, 1) if sample % 2 == 0 else (1, 0): + times[mode].append(latency(graphs[mode], 20, 16)) + result = dict( + m=m, + scope=scope, + median_us=list(map(statistics.median, times)), + samples_us=times, + bitwise_copy_replays=4, + final_projection_bitwise=True, + ) + rows.append(result) + print(json.dumps(result), flush=True) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps( + dict( + model=str(args.model), + layer=0, + tp4_rank=0, + synthetic_activations=True, + weight_copies=16, + torch=torch.__version__, + cuda=torch.version.cuda, + results=rows, + ), + indent=2, + ) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_hc_arithmetic_isolation.py b/benchmarks/kernels/benchmark_sm70_hc_arithmetic_isolation.py new file mode 100644 index 0000000000..1ae306477a --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_hc_arithmetic_isolation.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Isolate HC projection association using a retained failed chain input. + +Replays actual weights and all 96 HC modules without IPC or collectives. +Attention/MoE outputs remain fixed external inputs. This is a numerical +diagnostic, not distributed performance, model quality, or a runtime policy. +""" + +import argparse +import json +from pathlib import Path + +import torch +from safetensors import safe_open + +from benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv import check_exclusive +from benchmarks.kernels.benchmark_sm70_hc_tp4 import load_weights +from vllm.models.qwen4_exp.nvidia.ops.hc import ( + grouped_gemma_rmsnorm, + hc_combine, + hc_combine_norm, + hc_gate_mix, + hc_silu, +) + + +def difference(actual, expected): + delta = actual.float() - expected.float() + limit = 3e-3 + 3e-3 * expected.float().abs() + return { + "changed": int((actual != expected).sum()), + "outside_envelope": int((delta.abs() > limit).sum()), + "max_abs": float(delta.abs().max()), + "relative_l2": float(delta.norm() / expected.float().norm().clamp_min(1e-30)), + "finite": bool(torch.isfinite(actual).all()), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--failure", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--quality-inputs", type=int, default=16) + args = parser.parse_args() + if args.out.exists() or args.quality_inputs < 1: + parser.error("Require a fresh output and positive quality-input count") + check_exclusive() + torch.set_num_threads(4) + retained = torch.load(args.failure, weights_only=True, map_location="cpu") + initial = retained["initial"].cuda() + cores = retained["external_core_outputs"].cuda() + weights = load_weights(args.model) + assert len(weights) == len(cores) == 96 + mapping = json.loads((args.model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + norms = [] + for i in range(96): + role = ("attn", "mlp")[i % 2] + prefix = f"model.language_model.layers.{i // 2}.{role}_hyper_connection." + name = prefix + "hc_norm.weight" + with safe_open(args.model / mapping[name], framework="pt", device="cpu") as f: + norms.append(f.get_tensor(name).half().cuda()) + up_shards = [ + [ + up.view(4, 2560, 320)[:, r * 640 : (r + 1) * 640] + .reshape(2560, 320) + .contiguous() + for r in range(4) + ] + for _, up in weights + ] + + def project(x, i, split_down, split_up): + down, up = weights[i] + if split_down: + pieces = [ + torch.nn.functional.linear(x, down[r * 80 : r * 80 + 88]) + for r in range(4) + ] + raw = torch.cat([p[:, :80] for p in pieces], dim=-1) + injection = pieces[3][:, 80:84] + else: + packed = torch.nn.functional.linear(x, down) + raw, injection = packed[:, :320], packed[:, 320:324] + lora = hc_silu(raw, 4) + if split_up: + pieces = [ + torch.nn.functional.linear(lora, w).view(len(x), 4, 640) + for w in up_shards[i] + ] + gate = torch.cat(pieces, dim=-1).reshape(len(x), 10240) + else: + gate = torch.nn.functional.linear(lora, up) + return hc_gate_mix(x, gate, 4), injection + + def chain(split_down, split_up): + collected = [] + state, injection = initial, None + for i in range(96): + if i == 2: + state = hc_combine(state, cores[i - 1], injection, 4) + if i in (0, 2): + xn = grouped_gemma_rmsnorm(state, norms[i], 1e-6, 4) + else: + state, xn = hc_combine_norm( + state, cores[i - 1], injection, norms[i], 1e-6, 4 + ) + block, injection = project(xn, i, split_down, split_up) + collected.append((state, xn, block, injection)) + return collected + + modes = { + "reference": (False, False), + "both_sharded": (True, True), + "only_up_sharded": (False, True), + "only_down_sharded": (True, False), + } + graphs, outputs = {}, {} + for name, split in modes.items(): + for _ in range(2): + chain(*split) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs[name] = chain(*split) + graphs[name] = graph + + results = [] + generator = torch.Generator(device="cuda").manual_seed(20260905) + for sample in range(args.quality_inputs): + check_exclusive() + if sample: + initial.normal_(generator=generator) + cores.normal_(generator=generator) + for graph in graphs.values(): + graph.replay() + torch.accelerator.synchronize() + if sample == 0: + index = retained["hc_module"] + # A new ablation must first reproduce the retained negative result. + assert torch.equal( + outputs["reference"][index][1].cpu(), retained["reference"] + ) + assert torch.equal( + outputs["both_sharded"][index][1].cpu(), retained["candidate"] + ) + record = {"input": sample, "arms": {}} + for name in modes: + if name == "reference": + continue + first_failure = None + changed = 0 + max_error = 0.0 + for i, (a, b) in enumerate(zip(outputs[name], outputs["reference"])): + for field, actual, expected in zip( + ("state", "normalized", "block", "injection"), a, b + ): + stats = difference(actual, expected) + changed += stats["changed"] + max_error = max(max_error, stats["max_abs"]) + if first_failure is None and ( + not stats["finite"] or stats["outside_envelope"] + ): + first_failure = {"module": i, "field": field, **stats} + record["arms"][name] = { + "first_failure": first_failure, + "total_changed": changed, + "max_abs": max_error, + } + results.append(record) + print(json.dumps(record), flush=True) + args.out.write_text( + json.dumps( + { + "scope": "96-module arithmetic isolation, no IPC/performance admission", + "rows": len(initial), + "saved_four_rank_failure_reproduced": True, + "envelope": {"atol": 3e-3, "rtol": 3e-3}, + "inputs": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_hc_batch_tp4.py b/benchmarks/kernels/benchmark_sm70_hc_batch_tp4.py index abc007dad2..2e854cfe03 100644 --- a/benchmarks/kernels/benchmark_sm70_hc_batch_tp4.py +++ b/benchmarks/kernels/benchmark_sm70_hc_batch_tp4.py @@ -51,13 +51,46 @@ def mix_scatter(gate, x, out, RANK: tl.constexpr): tl.store(out + row * 2560 + col, acc / 4) +def shard_weights(down, up, rank, layout): + """Screen zero-copy output shards; do not change production weight loading.""" + up_view = up.view(4, 2560, 320)[:, rank * 640 : (rank + 1) * 640] + if layout == "views": + # Extra eight down rows are unused except rank3's four injections. + # The kernel already discards other ranks' col80..87 outputs. + return down.narrow(0, rank * 80, 88), up_view + if layout == "down-view": + local_down = down.narrow(0, rank * 80, 88) + else: + local_down = down.new_zeros(88, 10240) + local_down[:80].copy_(down[rank * 80 : (rank + 1) * 80]) + if rank == 3: + local_down[80:84].copy_(down[320:324]) + return local_down, up_view.reshape(2560, 320).contiguous() + + +def up_projection(lora, weight, output): + if weight.ndim == 2: + return torch.nn.functional.linear(lora, weight) + # Each branch writes disjoint columns of each output row. The batch + # dimension shares the lora input (stride zero), never the output cells. + torch.bmm( + lora.unsqueeze(0).expand(4, -1, -1), + weight.transpose(1, 2), + out=output.view(lora.shape[0], 4, 640).transpose(0, 1), + ) + return output + + def main(): p = argparse.ArgumentParser(description=__doc__) p.add_argument("--model", type=Path, required=True) p.add_argument("--tokens", default="4,8,16") p.add_argument("--layer", type=int, default=0) p.add_argument("--pair", choices=("attn", "mlp"), default="attn") - p.add_argument("--mode", choices=("full", "down"), default="full") + p.add_argument("--mode", choices=("full", "down", "fused"), default="full") + p.add_argument( + "--weight-layout", choices=("packed", "down-view", "views"), default="packed" + ) p.add_argument( "--weight-copies", type=int, @@ -68,6 +101,13 @@ def main(): "--profile", action="store_true", help="Capture graph nodes, skip timing sweep" ) p.add_argument("--out", type=Path, required=True) + p.add_argument("--check-only", action="store_true", help="Skip performance timing") + p.add_argument( + "--aux-check", + type=int, + default=0, + help="Fused mode only: interleave this many auxiliary sum2 graph replays", + ) args = p.parse_args() if args.weight_copies <= 0: p.error("--weight-copies must be positive") @@ -111,6 +151,18 @@ def check_exclusive(): ca = CustomAllreduce(group, rank, max_size=128 * 1024) assert not ca.disabled + hc_ca = None + if args.mode == "fused": + assert hasattr(torch.ops._C_custom_ar_flashnext, "hc_down_gather"), ( + "Build sm70_hc_push_gather.cuh into the SAME custom-AR sidecar " + "that owns this communicator" + ) + # Reuse the existing communicator implementation, but give HC a + # separate packet/epoch channel from auxiliary-stream collectives. + hc_ca = CustomAllreduce(group, rank, max_size=128 * 1024) + assert not hc_ca.disabled + elif args.aux_check: + p.error("--aux-check requires --mode fused") index = json.loads((args.model / "model.safetensors.index.json").read_text())[ "weight_map" ] @@ -125,18 +177,10 @@ def weight(suffix): injection = weight(".block_inject_weight.weight") down = torch.cat((down, injection, down.new_zeros(12, 10240))) up = weight(".input_mix_weight_up.weight") - local_down_w = down.new_zeros(88, 10240) - local_down_w[:80].copy_(down[rank * 80 : (rank + 1) * 80]) - if rank == 3: - local_down_w[80:84].copy_(down[320:324]) - local_up_w = ( - up.view(4, 2560, 320)[:, rank * 640 : (rank + 1) * 640] - .reshape(2560, 320) - .contiguous() - ) - weights = [(down, up, local_down_w, local_up_w)] + weights = [(down, up, *shard_weights(down, up, rank, args.weight_layout))] for _ in range(args.weight_copies - 1): - weights.append(tuple(w.clone() for w in weights[0])) + wd, wu = down.clone(), up.clone() + weights.append((wd, wu, *shard_weights(wd, wu, rank, args.weight_layout))) results = [] try: for m in map(int, args.tokens.split(",")): @@ -146,6 +190,9 @@ def weight(suffix): full_down = torch.empty_like(sparse_down) sparse_output = x.new_empty(m, 2560) output = torch.empty_like(sparse_output) + fused_lora = x.new_empty(m, 320) + fused_injection = x.new_empty(m, 4) + local_gate = x.new_empty(m, 2560) def baseline(index=0): wd, wu, _, _ = weights[index] @@ -155,9 +202,18 @@ def baseline(index=0): :, 320:324 ] - def candidate(registered, index=0): + def candidate(registered, index=0, use_fused=True): _, wu, local_wd, local_wu = weights[index] d = torch.nn.functional.linear(x, local_wd) + if args.mode == "fused" and use_fused: + torch.ops._C_custom_ar_flashnext.hc_down_gather( + hc_ca._ptr, d, fused_injection, fused_lora + ) + gate = up_projection(fused_lora, local_wu, local_gate) + torch.ops._C_custom_ar_flashnext.hc_mix_gather( + hc_ca._ptr, gate, x, output + ) + return output, fused_injection scatter_down[(m,)](d, sparse_down, RANK=rank) ca.all_reduce(sparse_down, out=full_down, registered=registered) lora = hc_silu(full_down[:, :320], 4) @@ -165,7 +221,7 @@ def candidate(registered, index=0): return hc_gate_mix( x, torch.nn.functional.linear(lora, wu), 4 ), full_down[:, 320:324] - gate = torch.nn.functional.linear(lora, local_wu) + gate = up_projection(lora, local_wu, local_gate) mix_scatter[(m, 10)](gate, x, sparse_output, RANK=rank) ca.all_reduce(sparse_output, out=output, registered=registered) return output, full_down[:, 320:324] @@ -191,8 +247,16 @@ def candidate(registered, index=0): dist.broadcast(x, 0) expected = tuple(t.clone() for t in baseline()) actual = tuple(t.clone() for t in candidate(False)) + if args.mode == "fused": + unfused = candidate(False, use_fused=False) + # Compare the fused publication separately from the local + # GEMM's already documented difference versus replicated. + for a, b in zip(actual, unfused): + torch.testing.assert_close(a, b, rtol=0, atol=0) sparse_down.fill_(float("nan")) output.fill_(float("nan")) + fused_lora.fill_(float("nan")) + fused_injection.fill_(float("nan")) graphs[1].replay() torch.cuda.synchronize() row = [] @@ -211,6 +275,53 @@ def candidate(registered, index=0): ) checks.append(row) assert all(c["finite"] for row in checks for c in row) + if args.aux_check: + performance_input = x.clone() + aux_input = x.new_full((m, 2560), (rank + 1) / 64) + aux_output = torch.empty_like(aux_input) + aux_stream = torch.cuda.Stream() + aux_graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + dist.barrier() + with ca.capture(), torch.cuda.graph(aux_graph, stream=aux_stream): + for _ in range(16): + ca.all_reduce_sum2(aux_input, aux_input, out=aux_output) + for cycle in range(args.aux_check): + # Distinct epochs/data, with deliberately skewed stream + # enqueue order across ranks. Never run unfused HC's ca + # collectives concurrently with this same auxiliary ca. + x.normal_().mul_(0.25 + cycle % 3) + dist.broadcast(x, 0) + expected_hc = tuple( + t.clone() for t in candidate(False, use_fused=False) + ) + aux_input.fill_((rank + 1 + cycle % 13) / 64) + aux_output.fill_(float("nan")) + output.fill_(float("nan")) + fused_lora.fill_(float("nan")) + fused_injection.fill_(float("nan")) + torch.cuda.synchronize() + dist.barrier() + if (rank + cycle) % 2: + with torch.cuda.stream(aux_stream): + aux_graph.replay() + graphs[1].replay() + else: + graphs[1].replay() + with torch.cuda.stream(aux_stream): + aux_graph.replay() + torch.cuda.synchronize() + for a, b in zip(outputs[1], expected_hc): + torch.testing.assert_close(a, b, rtol=0, atol=0) + expected_aux = (20 + 8 * (cycle % 13)) / 64 + torch.testing.assert_close( + aux_output, + torch.full_like(aux_output, expected_aux), + rtol=0, + atol=0, + ) + x.copy_(performance_input) + del aux_graph, aux_stream, aux_input, aux_output, performance_input if args.profile: for _ in range(30): graphs[0].replay() @@ -236,7 +347,7 @@ def candidate(registered, index=0): del graphs, outputs continue times = [[], []] - for sample in range(5): + for sample in range(0 if args.check_only else 5): for which in (0, 1) if sample % 2 == 0 else (1, 0): for _ in range(20): graphs[which].replay() @@ -259,10 +370,13 @@ def candidate(registered, index=0): local_result = dict( rank=rank, tokens=m, - baseline_us=statistics.median(times[0]), - candidate_us=statistics.median(times[1]), + baseline_us=statistics.median(times[0]) if times[0] else None, + candidate_us=statistics.median(times[1]) if times[1] else None, times=times, checks=checks, + fused_postops_zero_tolerance_check_passed=args.mode == "fused", + auxiliary_interleaved_replays=args.aux_check, + auxiliary_hc_inputs_changed=bool(args.aux_check), ) gathered = [None] * 4 dist.all_gather_object(gathered, local_result, group=group) @@ -278,12 +392,20 @@ def candidate(registered, index=0): "VLLM_SM70_TP4_PUSH_ALLREDUCE_SMALL_MESSAGES", "0" ), "calls_per_graph": calls_per_graph, + "private_hc_channel": hc_ca is not None, "replicated_weight_bytes": sum( w.numel() * w.element_size() for ws in weights for w in ws[:2] ), "sharded_weight_bytes": sum( w.numel() * w.element_size() for ws in weights for w in ws[2:] ), + "additional_sharded_storage_bytes": sum( + shard.numel() * shard.element_size() + for ws in weights + for original, shard in zip(ws[:2], ws[2:]) + if original.untyped_storage().data_ptr() + != shard.untyped_storage().data_ptr() + ), } if library: runtime["custom_ar_sha256"] = hashlib.sha256( @@ -299,6 +421,8 @@ def candidate(registered, index=0): + "\n" ) finally: + if hc_ca is not None: + hc_ca.close() ca.close() dist.destroy_process_group(group) dist.destroy_process_group() diff --git a/benchmarks/kernels/benchmark_sm70_hc_full_chain.py b/benchmarks/kernels/benchmark_sm70_hc_full_chain.py index 2b29c0f784..00532acb61 100644 --- a/benchmarks/kernels/benchmark_sm70_hc_full_chain.py +++ b/benchmarks/kernels/benchmark_sm70_hc_full_chain.py @@ -18,6 +18,7 @@ import json import os import subprocess +from contextlib import nullcontext from pathlib import Path from statistics import median from types import SimpleNamespace @@ -37,12 +38,14 @@ hc_gate_mix, hc_silu, ) +from vllm.models.qwen4_exp.nvidia.sm70_batch_hc import _batch_hc def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--model", type=Path, required=True) parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--rows", type=int, choices=(1, 4, 8, 16), default=1) parser.add_argument("--quality-inputs", type=int, default=16) parser.add_argument("--warmup", type=int, default=1000) parser.add_argument("--replays", type=int, default=150) @@ -58,6 +61,12 @@ def main() -> None: help="Auxiliary sum2 replays per changing input with --fused-up", ) args = parser.parse_args() + if args.out.exists(): + parser.error("Refusing to overwrite previous evidence") + if args.rows > 1 and args.fused_up: + parser.error( + "--fused-up is the original M1 comparison; batch compares full chains" + ) if args.fused_up and not envs.VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1: raise RuntimeError( "Set VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1=1 for the aux gate" @@ -72,7 +81,14 @@ def main() -> None: dist.init_process_group("nccl") group = dist.new_group(backend="gloo") comm = CustomAllreduce(group=group, device=local_rank, max_size=8 * 1024 * 1024) + batch_comm = None try: + if args.rows > 1: + batch_comm = CustomAllreduce( + group=group, device=local_rank, max_size=128 * 1024 + ) + if not batch_comm.supports_sm70_qwen38_hc_batch(): + raise RuntimeError("Load the source-matched batch HC communicator") owned_pids = [None] * 4 dist.all_gather_object(owned_pids, os.getpid(), group=group) @@ -125,14 +141,30 @@ def get(name: str) -> torch.Tensor: final_up = get(prefix + "hyper_connection_mixer.input_mix_weight_up.weight") gen = torch.Generator(device="cuda").manual_seed(20260905) initial = torch.randn( - (1, 10240), device="cuda", dtype=torch.float16, generator=gen + (args.rows, 10240), device="cuda", dtype=torch.float16, generator=gen ) cores = torch.randn( - (96, 1, 2560), device="cuda", dtype=torch.float16, generator=gen + (96, args.rows, 2560), device="cuda", dtype=torch.float16, generator=gen ) - if not comm.can_sm70_qwen38_hc_shard(initial): + if args.rows == 1 and not comm.can_sm70_qwen38_hc_shard(initial): raise RuntimeError("The exact TP4 HC route is unavailable") - tp = SimpleNamespace(device_communicator=SimpleNamespace(ca_comm=comm)) + if batch_comm is not None and not batch_comm.can_sm70_qwen38_hc_batch(initial): + raise RuntimeError("Batch HC does not support the requested rows") + tp = SimpleNamespace( + device_communicator=SimpleNamespace( + ca_comm=comm, sm70_hc_batch_comm=batch_comm + ) + ) + up_shards = ( + [ + up.view(4, 2560, 320)[:, rank * 640 : (rank + 1) * 640] + .reshape(2560, 320) + .contiguous() + for _, up in weights + ] + if args.rows > 1 + else [] + ) if args.fused_up: sum_gen = torch.Generator(device="cuda").manual_seed(20260905 + rank) sum_a = torch.randn( @@ -157,7 +189,7 @@ def finish(state: torch.Tensor, injection: torch.Tensor): return combined, hc_gate_mix(xn, gate, 4) # A model's normal warmup initializes cuBLAS before graph capture. - finish(initial, torch.zeros((1, 4), device="cuda", dtype=torch.float16)) + finish(initial, torch.zeros((args.rows, 4), device="cuda", dtype=torch.float16)) torch.cuda.synchronize() def capture(mode: str, overlap: bool = False): @@ -179,7 +211,12 @@ def capture(mode: str, overlap: bool = False): "supports_sm70_qwen38_hc_up_mix_allgather", return_value=mode == "fused", ), + patch( + "vllm.models.qwen4_exp.nvidia.sm70_batch_hc._decode_context_ok", + return_value=mode == "batch", + ), comm.capture(), + batch_comm.capture() if batch_comm is not None else nullcontext(), torch.cuda.graph(graph), ): main_stream = torch.cuda.current_stream() @@ -197,9 +234,14 @@ def capture(mode: str, overlap: bool = False): state, xn = hc_combine_norm( state, cores[i - 1], injection, norms[i], 1e-6, 4 ) - block, injection = torch.ops.vllm.qwen38_sm70_fp16_fused_hc( - xn, down, up - ) + if args.rows > 1: + block, injection = _batch_hc( + xn, down, up, up_shards[i], True, False + ) + else: + block, injection = torch.ops.vllm.qwen38_sm70_fp16_fused_hc( + xn, down, up + ) outputs.extend((state, xn, block, injection)) if aux is not None: with torch.cuda.stream(aux): @@ -211,7 +253,13 @@ def capture(mode: str, overlap: bool = False): dist.barrier() return graph, outputs, sums - timed_modes = ("hidden", "fused") if args.fused_up else ("gate", "hidden") + timed_modes = ( + ("legacy", "batch") + if args.rows > 1 + else ("hidden", "fused") + if args.fused_up + else ("gate", "hidden") + ) graphs = {mode: capture(mode) for mode in timed_modes} if args.fused_up: graphs["fused_aux"] = capture("fused", overlap=True) @@ -236,6 +284,41 @@ def replay_and_check(stress: bool): [x.flatten().view(torch.int16) for x in graphs[mode][1]] ) diffs[mode] = int(torch.count_nonzero(expected != actual)) + if args.rows > 1: + # Preserve the original M1 bitwise gate. The batched + # sharded GEMMs have independent rounding; bound every + # intermediate instead of silently ignoring differences. + for index, (a, b) in enumerate( + zip(graphs[mode][1], graphs[timed_modes[0]][1], strict=True) + ): + try: + torch.testing.assert_close(a, b, atol=3e-3, rtol=3e-3) + except AssertionError: + failure = args.out.with_suffix(f".rank{rank}.failure.pt") + torch.save( + { + "scope": "HC numerical gate failure, not speed", + "mode": mode, + "rows": args.rows, + "output_index": index, + "hc_module": index // 4, + "component": ( + "state", + "normalized", + "block", + "injection", + )[index % 4] + if index < 384 + else "final_mixer", + "candidate": a.cpu(), + "reference": b.cpu(), + "initial": initial.cpu(), + "external_core_outputs": cores.cpu(), + }, + failure, + ) + print(f"HC numerical gate failed: {failure}", flush=True) + raise sum_diff = 0 if args.fused_up: actual_sum = torch.stack(graphs["fused_aux"][2]) @@ -272,7 +355,7 @@ def replay_and_check(stress: bool): ) if rank == 0: print({"quality": quality}, flush=True) - if any(q["mismatches"] for q in quality): + if args.rows == 1 and any(q["mismatches"] for q in quality): raise RuntimeError("Full HC outputs are not bitwise") ensure_exclusive() for mode in timed_modes: @@ -282,7 +365,7 @@ def replay_and_check(stress: bool): torch.cuda.synchronize() dist.barrier() samples = {mode: [] for mode in timed_modes} - for repeat in range(3): + for repeat in range(5): modes = timed_modes if repeat % 2 == 0 else timed_modes[::-1] for mode in modes: ensure_exclusive() @@ -313,7 +396,7 @@ def replay_and_check(stress: bool): {"rank": rank, "hc_mismatches": diffs, "sum2_mismatches": sum_diff}, group=group, ) - if any( + if args.rows == 1 and any( any(q["hc_mismatches"].values()) or q["sum2_mismatches"] for q in post_quality ): @@ -340,6 +423,10 @@ def replay_and_check(stress: bool): "gpu": torch.cuda.get_device_name(), "visible_devices": visible, "quality": quality, + "rows": args.rows, + "numerical_gate": "bitwise" + if args.rows == 1 + else "every intermediate atol=rtol=3e-3", "quality_inputs": args.quality_inputs, "post_timing_quality": post_quality, "aux_stress_replays": args.quality_inputs * args.aux_stress_replays @@ -351,6 +438,8 @@ def replay_and_check(stress: bool): args.out.write_text(json.dumps(result, indent=2) + "\n") print(json.dumps(result, indent=2), flush=True) finally: + if batch_comm is not None: + batch_comm.close() comm.close() dist.destroy_process_group(group) dist.destroy_process_group() diff --git a/benchmarks/kernels/benchmark_sm70_moe_w2_locality.py b/benchmarks/kernels/benchmark_sm70_moe_w2_locality.py new file mode 100644 index 0000000000..ef1af4edc0 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_moe_w2_locality.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: B023 +"""Exact-arithmetic MoE locality screens using real weights and route traces. + +Screen W2 warp mapping, paired W13 projections, or tile-major scale layout. +Measure the affected component and complete grouped MoE, including planning, +scatter and reduction. Activations are synthetic; this is not a quality gate. +""" + +import argparse +import glob +import hashlib +import json +import statistics +from pathlib import Path + +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv import check_exclusive +from benchmarks.kernels.benchmark_sm70_moe_packed_w13 import ( + checkpoint_weights, + graph, + latency, +) +from benchmarks.kernels.sm70_paired_stats import paired_latency_interval + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--library", type=Path, required=True) + experiment = parser.add_mutually_exclusive_group() + experiment.add_argument("--w13-pair", action="store_true") + experiment.add_argument("--w13-single-tile", action="store_true") + experiment.add_argument("--scale-layout", action="store_true") + experiment.add_argument("--compact-tasks", action="store_true") + parser.add_argument("--compact-w2-only", action="store_true") + parser.add_argument("--routes", required=True, help="Glob of saved [M,10] tensors") + parser.add_argument("--route-limit", type=int, default=1) + parser.add_argument("--tokens", type=int, nargs="+", default=[4, 8, 16]) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--samples", type=int, default=5) + args = parser.parse_args() + if args.out.exists(): + parser.error("Refusing to overwrite a previous result") + if any(m not in (4, 8, 16) for m in args.tokens): + parser.error("This paired screen supports M4/8/16") + if min(args.route_limit, args.repeats, args.samples) < 1: + parser.error("Counts must be positive") + if args.compact_tasks and args.samples != 5: + parser.error("Compact-task admission uses five paired timing blocks") + if args.compact_w2_only and not args.compact_tasks: + parser.error("--compact-w2-only requires --compact-tasks") + assert torch.cuda.get_device_capability() == (7, 0) + check_exclusive() + torch.ops.load_library(str(args.library.resolve())) + native = torch.ops._C + w13_variant = args.w13_pair or args.w13_single_tile + candidate = ( + torch.ops._C_moe_compact_tasks + if args.compact_tasks + else torch.ops._C_moe_scale_layout + if args.scale_layout + else torch.ops._C_moe_single_tile + if args.w13_single_tile + else torch.ops._C_moe_pair + if args.w13_pair + else torch.ops._C_moe_locality + ) + paths = sorted(map(Path, glob.glob(args.routes)))[: args.route_limit] + if not paths: + parser.error("No route files matched") + captured = {} + for path in paths: + value = torch.load(path, map_location="cpu", weights_only=True) + value = value["tensor"] if isinstance(value, dict) else value + assert value.ndim == 2 and value.shape[1] == 10 + assert value.shape[0] >= max(args.tokens) + captured[path.name] = value.to(torch.int32) + torch.manual_seed(20260905) + w13, s13, w2, s2 = checkpoint_weights(args.model, 0, 0, True) + resources = ( + {} + if args.scale_layout or args.compact_tasks + else { + str(mode): list(candidate.resources(mode)) + for mode in ((4, 5, 8) if w13_variant else (1, 2, 4)) + } + ) + if args.scale_layout: + tile_s13 = s13.reshape(512, 160, 10, 32).permute(0, 2, 1, 3).contiguous() + tile_s2 = s2.reshape(512, 10, 80, 32).permute(0, 2, 1, 3).contiguous() + results = [] + for m in args.tokens: + n = m * 10 + split = {4: 5, 8: 4, 16: 8}[m] + x = torch.randn(m, 2560, device="cuda", dtype=torch.float16) * 0.1 + ids = torch.empty(n, device="cuda", dtype=torch.int32) + topk = torch.softmax(torch.randn(m, 10, device="cuda"), -1) + mid = torch.empty(n, 160, device="cuda", dtype=torch.float16) + routed = torch.empty(n, 2560, device="cuda", dtype=torch.float16) + rows = torch.empty(n, 8, device="cuda", dtype=torch.int32) + experts = torch.empty(n, device="cuda", dtype=torch.int32) + sizes = torch.empty_like(experts) + total = torch.empty(1, device="cuda", dtype=torch.int32) + modes = ( + (0, 2) + if args.compact_w2_only + else (0, 1, 2, 3) + if args.scale_layout or args.compact_tasks + else (0, 1) + if w13_variant + else (0, 1, 2, 4) + ) + outputs = {mode: torch.empty_like(x) for mode in modes} + + def w13_call(mode=0): + if args.compact_tasks and mode & 1: + candidate.nvfp4_grouped_w13_sm70_out( + mid, x, w13, s13, ids, rows, experts, sizes, total, split, True + ) + return + if args.scale_layout and mode & 1: + candidate.nvfp4_grouped_w13_sm70_out( + mid, x, w13, tile_s13, ids, rows, experts, sizes, total, split, True + ) + return + run = ( + candidate.run + if w13_variant and mode + else native.nvfp4_grouped_w13_sm70_out + ) + run(mid, x, w13, s13, ids, rows, experts, sizes, total, split, True) + + def w2_call(mode): + params = ( + outputs[mode], + routed, + mid, + w2, + tile_s2 if args.scale_layout and mode & 2 else s2, + topk, + rows, + experts, + sizes, + total, + ) + if args.scale_layout and mode & 2 or args.compact_tasks and mode & 2: + candidate.nvfp4_grouped_w2_sm70_out(*params) + elif mode == 0 or w13_variant or args.scale_layout or args.compact_tasks: + native.nvfp4_grouped_w2_sm70_out(*params) + else: + candidate.w2(*params, mode) + + def complete(mode): + w13_call(mode) + w2_call(mode) + + cases = { + "distinct": torch.arange(n).reshape(m, 10), + "shared10": torch.arange(10).repeat(m, 1), + **{name: value[:m] for name, value in captured.items()}, + } + for name, values in cases.items(): + ids.copy_(values.reshape(-1)) + full = {mode: graph(lambda: complete(mode)) for mode in outputs} + changed_cases = ( + values, + values.flip(0), + torch.arange(10).repeat(m, 1), + torch.zeros(m, 10, dtype=torch.int32), + torch.arange(n).reshape(m, 10), + torch.full((m, 10), -1), + torch.full((m, 10), 512), + ) + for changed in changed_cases: + x.normal_(0, 0.1) + topk.copy_(torch.softmax(torch.randn_like(topk), -1)) + ids.copy_(changed.reshape(-1)) + for mode in outputs: + for buf in (rows, experts, sizes, total): + buf.fill_(-12345) + mid.fill_(float("nan")) + routed.fill_(float("nan")) + outputs[mode].fill_(float("nan")) + full[mode].replay() + if mode == 0: + reference_mid = mid.clone() + assert torch.equal(mid, reference_mid), (m, name, mode, "W13") + assert torch.isfinite(outputs[mode]).all() + assert torch.equal(outputs[mode], outputs[0]), (m, name, mode) + ids.copy_(values.reshape(-1)) + w13_call() + count = total.item() + parts = ( + (("w2", w2_call),) + if args.compact_w2_only + else (("w13", w13_call), ("w2", w2_call)) + if args.scale_layout or args.compact_tasks + else (("w13", w13_call),) + if w13_variant + else (("w2", w2_call),) + ) + scopes = [ + (name, {mode: graph(lambda: part(mode)) for mode in outputs}) + for name, part in parts + ] + scopes.append(("complete_moe", full)) + for scope, graphs in scopes: + times = {mode: [] for mode in outputs} + for sample in range(args.samples): + order = modes if sample % 2 == 0 else tuple(reversed(modes)) + for mode in order: + check_exclusive() + times[mode].append(latency(graphs[mode], args.repeats)) + result = { + "m": m, + "case": name, + "scope": scope, + "groups": count, + "split": split, + "exact_changed_replays": len(changed_cases), + "median_us": { + str(k): statistics.median(v) for k, v in times.items() + }, + "samples_us": times, + "paired_intervals": { + str(mode): paired_latency_interval(times[0], times[mode]) + for mode in modes + if mode and args.samples == 5 + }, + } + results.append(result) + print(json.dumps(result), flush=True) + report = { + "model": str(args.model), + "layer": 0, + "tp4_rank": 0, + "synthetic_activations": True, + "w13_pair": args.w13_pair, + "w13_single_tile": args.w13_single_tile, + "scale_layout": args.scale_layout, + "compact_tasks": args.compact_tasks, + "compact_w2_only": args.compact_w2_only, + "graph_unroll": 16, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "library_sha256": hashlib.sha256(args.library.read_bytes()).hexdigest(), + "route_sha256": { + str(p): hashlib.sha256(p.read_bytes()).hexdigest() for p in paths + }, + "resource_fields": [ + "registers", + "shared_bytes", + "local_bytes", + "max_ctas_per_sm", + ], + "resources": resources, + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qsa_batch_routes.py b/benchmarks/kernels/benchmark_sm70_qsa_batch_routes.py new file mode 100644 index 0000000000..70399f8b1f --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qsa_batch_routes.py @@ -0,0 +1,323 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen QSA attention for independent no-MTP requests, not verifier rows. + +Compare forced Triton, direct XQA, and padded grouped Page4 with the same +metadata. This is attention only, not the indexer, model round or quality score. +""" + +import argparse +import hashlib +import json +import os +import statistics +import subprocess +from pathlib import Path + +import torch +from flash_attn_v100.flash_attn_interface import flash_attn_v100_cuda as native + +from vllm.models.qwen4_exp.nvidia.ops import qsa as ops + + +def check_exclusive(): + report = subprocess.check_output( + [ + "nvidia-smi", + "-i", + os.environ["CUDA_VISIBLE_DEVICES"], + "--query-compute-apps=pid,process_name", + "--format=csv,noheader", + ], + text=True, + ) + for line in report.splitlines(): + pid, _, name = line.partition(",") + if ( + pid.strip().isdigit() + and int(pid) != os.getpid() + and "snapd-desktop-integration" not in name + ): + raise RuntimeError(f"Foreign GPU process; discard timing: {line}") + + +@torch.inference_mode() +def screen(rows, args): + device = "cuda" + page = 784 + blocks = (args.seq_len + page - 1) // page + q = torch.randn((rows, 6, 256), device=device, dtype=torch.float16) + k = torch.randn((rows * blocks, page, 1, 256), device=device, dtype=torch.float16) + v = torch.randn_like(k) + generator = torch.Generator().manual_seed(7) + selected_blocks = torch.stack( + [ + torch.randperm(args.seq_len // 4 - 1, generator=generator)[:512] + for _ in range(rows) + ] + ).to(device=device, dtype=torch.int32) + # Distinct requests have distinct physical pages, even at equal positions. + table = torch.randperm(rows * blocks, device=device, dtype=torch.int32) + table = table.view(rows, blocks) + req = torch.arange(rows, device=device, dtype=torch.int32) + pos = torch.full((rows,), args.seq_len - 1, device=device, dtype=torch.int64) + lengths = torch.full((rows,), args.seq_len, device=device, dtype=torch.int32) + indices = ops.expand_qsa_block_indices_cuda( + selected_blocks, + pos, + lengths, + req, + 4, + 2048, + ) + outputs = [torch.empty_like(q) for _ in range(3)] + padded = ((rows + 7) // 8) * 8 + pq = q.new_zeros((padded, 6, 256)) + pi = indices.new_full((padded, indices.shape[1]), -1) + pr = req.new_full((padded,), -1) + pp = pos.new_zeros(padded) + po = torch.empty_like(pq) + + def triton_call(): + ops.qsa_sparse_paged_attention( + q, + k, + v, + indices, + table, + req, + out=outputs[0], + query_positions=pos, + sequence_lengths=lengths, + ) + + def xqa_call(): + ops._qsa_sparse_paged_attention_sm70_xqa_page4_batch( + q, + k, + v, + indices, + table, + req, + pos, + lengths, + outputs[1], + "auto", + 1.0, + 1.0, + native, + ) + + def grouped_call(): + pq[:rows].copy_(q) + pi[:rows].copy_(indices) + pr[:rows].copy_(req) + pp[:rows].copy_(pos) + ops._qsa_sparse_paged_attention_sm70_grouped_page4( + pq, + k, + v, + pi, + table, + pr, + pp, + lengths, + po, + "auto", + 1.0, + 1.0, + native, + ) + outputs[2].copy_(po[:rows]) + + calls = [triton_call, xqa_call, grouped_call] + graphs = [] + stream = torch.cuda.Stream() + for call in calls: + with torch.cuda.stream(stream): + for _ in range(3): + call() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for _ in range(args.calls_per_graph): + call() + graphs.append(graph) + + checks = [] + passed = [True] * 3 + for cycle in range(6): + q.normal_().mul_((0.25, 1.0, 3.0)[cycle % 3]) + # A canonical QSA tail has (visible_length % 4) entries, not always + # three. Sweep all residues while keeping captured pointers fixed. + lengths.fill_(args.seq_len - cycle % 4) + pos.copy_(lengths.long() - 1) + ops.expand_qsa_block_indices_cuda( + selected_blocks, + pos, + lengths, + req, + 4, + 2048, + out=indices, + ) + complete = ( + selected_blocks[:, :, None] * 4 + torch.arange(4, device=device) + ).reshape(rows, 2048) + tail_offset = torch.arange(3, device=device) + tail = (lengths[:, None] // 4) * 4 + tail_offset + tail = torch.where(tail_offset < lengths[:, None] % 4, tail, -1) + torch.testing.assert_close( + indices.long(), torch.cat((complete, tail), 1).long() + ) + # Exercise mutable slot mappings while graph pointers remain fixed. + table.copy_( + torch.randperm(rows * blocks, device=device, dtype=torch.int32).view_as( + table + ) + ) + safe = indices.clamp_min(0) + physical = table.gather(1, (safe // page).long()).long() + selected_k = k[physical, (safe % page).long(), 0].float() + selected_v = v[physical, (safe % page).long(), 0].float() + scores = q.float() @ selected_k.transpose(1, 2) / 16 + scores.masked_fill_((indices < 0)[:, None, :], float("-inf")) + expected = torch.softmax(scores, dim=-1) @ selected_v + for call in calls: + call() + eager = [out.clone() for out in outputs] + for out in outputs: + out.fill_(float("nan")) + po.fill_(float("nan")) + for graph in graphs: + graph.replay() + torch.cuda.synchronize() + errors = [] + for route, (out, ref) in enumerate(zip(outputs, eager)): + exact = torch.equal(out, ref) + close = torch.allclose(out.float(), expected, atol=2e-3, rtol=1e-2) + err = out.float() - expected + rel = (err.norm() / expected.norm()).item() + passed[route] &= exact and close and rel <= 5e-3 + errors.append( + { + "max_abs": err.abs().max().item(), + "relative_l2": rel, + "graph_equals_eager": exact, + "fp32_oracle_close": close, + } + ) + if torch.count_nonzero(po[rows:]).item(): + passed[2] = False + checks.append(errors) + + print( + json.dumps( + { + "rows": rows, + "route_order": ["triton", "direct_xqa", "padded_grouped"], + "micro_gate_passed": passed, + "checks": checks, + } + ), + flush=True, + ) + + # Timing retains the declared length, not the final residue-check length. + lengths.fill_(args.seq_len) + pos.copy_(lengths.long() - 1) + ops.expand_qsa_block_indices_cuda( + selected_blocks, + pos, + lengths, + req, + 4, + 2048, + out=indices, + ) + samples = [[] for _ in calls] + for sample in range(5): + check_exclusive() + for i in range(3) if sample % 2 == 0 else reversed(range(3)): + # Retain diagnostics, but never admit timing for a failed oracle. + if not passed[0] or not passed[i]: + continue + graph = graphs[i] + for _ in range(5): + graph.replay() + start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + start.record() + for _ in range(args.replays): + graph.replay() + end.record() + end.synchronize() + samples[i].append( + start.elapsed_time(end) * 1000 / args.replays / args.calls_per_graph + ) + check_exclusive() + return { + "rows": rows, + "independent_requests": True, + "query_heads": 6, + "kv_heads": 1, + "head_dim": 256, + "kv_dtype": "float16", + "selection_width": indices.shape[1], + "padded_rows": padded, + "samples_us": samples, + "median_us": [statistics.median(s) if s else None for s in samples], + "route_order": ["triton", "direct_xqa", "padded_grouped"], + "changed_input_slot_map_cycles": checks, + "micro_gate_passed": passed, + } + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--rows", default="1,4,8,16") + p.add_argument("--seq-len", type=int, default=8192) + p.add_argument("--replays", type=int, default=40) + p.add_argument("--calls-per-graph", type=int, default=16) + p.add_argument("--out", type=Path, required=True) + args = p.parse_args() + if args.seq_len < 4096 or args.seq_len % 4: + p.error("seq-len must be a multiple of four and at least 4096") + rows = [int(m) for m in args.rows.split(",")] + if not rows or min(rows) < 1 or min(args.replays, args.calls_per_graph) < 1: + p.error("rows and replay counts must be positive") + if torch.cuda.get_device_capability() != (7, 0): + raise RuntimeError("SM70 required") + if not ops._qsa_grouped_page4_supported(native, "auto"): + raise RuntimeError("grouped Page4 native capability unavailable") + original_gate = ops._SM70_QSA_XQA_PAGE4 + try: + # Force the reference; candidate calls bypass the gate explicitly. + ops._SM70_QSA_XQA_PAGE4 = False + torch.manual_seed(7) + results = [screen(m, args) for m in rows] + finally: + ops._SM70_QSA_XQA_PAGE4 = original_gate + payload = { + "qsa_source": ops.__file__, + "native": native.__file__, + "native_sha256": hashlib.sha256(Path(native.__file__).read_bytes()).hexdigest(), + "grouped_abi": ops._qsa_grouped_page4_abi_version(native), + "torch": torch.__version__, + "device": torch.cuda.get_device_name(), + "seq_len": args.seq_len, + "calls_per_graph": args.calls_per_graph, + "replays": args.replays, + "results": results, + "scope": "Single-GPU attention micro, not TP4/model quality or throughput", + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(payload, indent=2) + "\n") + print(json.dumps(payload, indent=2)) + if not all(all(r["micro_gate_passed"]) for r in results): + raise SystemExit( + "A QSA micro quality gate failed; inspect retained diagnostics" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qsa_rounding_isolation.py b/benchmarks/kernels/benchmark_sm70_qsa_rounding_isolation.py new file mode 100644 index 0000000000..af0a157a2e --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qsa_rounding_isolation.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU numerical counterfactual on retained real QSA inputs. + +All arithmetic is FP64 except the explicitly modeled probability and output +materializations. This isolates their effect without GPU acquisition; it is +not a CUDA-kernel emulator, model score, performance test, or release gate. +""" + +import argparse +import json +import math +from pathlib import Path + +import torch + + +def materialize(result, gate): + result = result.half() + if gate is not None: + result = (result.float() * gate.reshape_as(result).float().sigmoid()).half() + return result + + +def isolate(capture, rounded): + q, k, v = (capture[name] for name in ("q", "k", "v")) + assert k.shape[2] == 1, "This diagnostic models one KV head per rank" + indices, table, requests = ( + capture[name].long() for name in ("indices", "table", "requests") + ) + safe = indices.clamp_min(0) + valid = ( + (indices >= 0) + & (safe // k.shape[1] < table.shape[1]) + & (requests[:, None] >= 0) + & (requests[:, None] < table.shape[0]) + ) + pages = table[ + requests.clamp(0, table.shape[0] - 1)[:, None], + (safe // k.shape[1]).clamp_max(table.shape[1] - 1), + ] + valid &= (pages >= 0) & (pages < k.shape[0]) + keys = k[pages.clamp(0, k.shape[0] - 1), safe % k.shape[1], 0].double() + values = v[pages.clamp(0, v.shape[0] - 1), safe % v.shape[1], 0].double() + # Invalid values may be NaN; masking their probability alone is not enough. + keys = torch.where(valid[:, :, None], keys, 0.0) + values = torch.where(valid[:, :, None], values, 0.0) + scores = torch.bmm(q.double(), keys.transpose(1, 2)) / math.sqrt(q.shape[2]) + scores.masked_fill_(~valid[:, None], -1e20) + tiles = (indices.shape[1] + 15) // 16 + target = 64 if len(q) <= 8 else 32 + splits = min(1 << (tiles.bit_length() - 1), target) + partials, lses = [], [] + for split in range(splits): + maximum = torch.full(q.shape[:2], -1e20, dtype=torch.float64) + denominator = torch.zeros_like(maximum) + numerator = torch.zeros(q.shape, dtype=torch.float64) + for tile in range(split * tiles // splits, (split + 1) * tiles // splits): + start, end = tile * 16, min((tile + 1) * 16, indices.shape[1]) + logits = scores[:, :, start:end] + next_max = torch.maximum(maximum, logits.max(-1).values) + alpha = (maximum - next_max).exp() + p = torch.where( + valid[:, None, start:end], (logits - next_max[:, :, None]).exp(), 0 + ) + pv = p.half().double() if rounded else p + numerator = numerator * alpha[:, :, None] + torch.bmm( + pv, values[:, start:end] + ) + denominator = denominator * alpha + p.sum(-1) + maximum = next_max + partials.append(numerator / denominator.clamp_min(1e-20)[:, :, None]) + lses.append( + torch.where(denominator > 0, maximum + denominator.log(), -torch.inf) + ) + lse = torch.stack(lses) + weights = (lse - lse.max(0).values).exp().nan_to_num(0) + merged = (torch.stack(partials) * weights[:, :, :, None]).sum(0) + merged /= weights.sum(0).clamp_min(1e-20)[:, :, None] + return materialize(merged, capture["gate"]) + + +def error(a, b): + a, b = a.double(), b.double() + return { + "relative_l2": float((a - b).norm() / b.norm().clamp_min(1e-30)), + "max_abs": float((a - b).abs().max()), + "changed": int((a != b).sum()), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--captures", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if args.out.exists(): + parser.error("Use a fresh output path") + paths = sorted(args.captures.glob("qsa-r*-n*.pt")) + if not paths: + parser.error("No retained QSA inputs found") + torch.set_num_threads(4) + results = [] + for path in paths: + capture = torch.load(path, weights_only=True, map_location="cpu") + rounded, full = isolate(capture, True), isolate(capture, False) + record = { + "capture": path.name, + "rounded_vs_reference": error(rounded, capture["reference"]), + "full_vs_reference": error(full, capture["reference"]), + "rounded_vs_native": error(rounded, capture["candidate"]), + "full_vs_native": error(full, capture["candidate"]), + "full_vs_oracle": error(full, capture["oracle"]), + } + results.append(record) + print(json.dumps(record), flush=True) + args.out.write_text( + json.dumps( + { + "scope": "CPU FP64 counterfactual, not CUDA emulation or quality proof", + "cases": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qsa_scorer_grid.py b/benchmarks/kernels/benchmark_sm70_qsa_scorer_grid.py new file mode 100644 index 0000000000..a6a5068257 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qsa_scorer_grid.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen QSA scorer CTA amortization; no production dispatch change. + +Keep max-context capacity fixed while changing live lengths. Compare the +existing kernel's contiguous tile grouping against its production grouping, +including changed-length CUDA Graph replay. Timing is operator-only, not +end-to-end throughput or model-quality admission. +""" + +import argparse +import json +import math +import statistics +from pathlib import Path + +import torch + +from benchmarks.kernels.sm70_qsa_strided_scorer import strided_qsa_mqa_paged_kernel +from vllm.models.qwen4_exp.nvidia.ops.qsa import _qsa_mqa_paged_kernel +from vllm.triton_utils import triton + + +def run_case(rows, length, groups, repeats, strided_stages=2): + device = torch.device("cuda") + # FlashNext's indexer projection is replicated, NOT sharded over TP4. + heads, dim, ratio, page_size, table_width = 4, 128, 4, 196, 335 + columns = page_size * table_width + block_n = 32 if rows == 1 else 64 + torch.manual_seed(20260905) + q = torch.randn(rows, heads, dim, dtype=torch.float16, device=device) + cache = torch.randn( + rows * table_width, page_size, 1, dim, dtype=torch.float16, device=device + ) + table = torch.randperm(rows * table_width, device=device).to(torch.int32) + table = table.reshape(rows, table_width) + request = torch.arange(rows, dtype=torch.int32, device=device) + positions = torch.full((rows,), length - 1, dtype=torch.int32, device=device) + lengths = torch.full_like(positions, length) + logits = { + group: torch.full((rows, columns), float("nan"), device=device) + for group in groups + } + visible = {group: torch.empty_like(positions) for group in groups} + + def launch(group): + out = logits[group] + strided = group < 0 + grid = ( + min(-group, triton.cdiv(columns, block_n)) + if strided + else triton.cdiv(columns, block_n * group) + ) + kernel = strided_qsa_mqa_paged_kernel if strided else _qsa_mqa_paged_kernel + return kernel[(rows, grid)]( + q, + cache, + table, + request, + positions, + lengths, + visible[group], + out, + *q.stride(), + cache.stride(0), + cache.stride(1), + cache.stride(3), + *table.stride(), + out.stride(0), + rows, + columns, + cache.shape[0], + rows, + math.sqrt(dim), + PAGE_SIZE=page_size, + PAGE_TABLE_WIDTH=table_width, + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_N=block_n, + BLOCK_D=dim, + TILES_PER_PROG=1 if strided else group, + STAGES=strided_stages if strided else 2, + MAX_N=16, + COMPRESS_RATIO=ratio, + num_warps=2, + ) + + graphs = {} + resources = {} + for group in groups: + for _ in range(3): + compiled = launch(group) + resources[group] = { + "shared_bytes": compiled.metadata.shared, + "registers_per_thread": compiled.n_regs, + "spills": compiled.n_spills, + } + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(group) + graphs[group] = graph + + # Shrink, grow, mixed lengths and page residues must reuse the same graph. + for iteration in range(4): + values = [max(1, length - (iteration * 197 + row * 31)) for row in range(rows)] + if iteration == 3: + values = [262144 - row * 3 for row in range(rows)] + lengths.copy_(torch.tensor(values, dtype=torch.int32, device=device)) + positions.copy_(lengths - 1) + q.mul_(-0.875) + for group in groups: + logits[group].fill_(float("nan")) + graphs[group].replay() + assert torch.equal(visible[1], lengths // ratio) + valid = torch.arange(columns, device=device)[None, :] < visible[1][:, None] + reference = logits[1][valid] + assert torch.isfinite(reference).all() + for group in groups: + assert torch.equal(visible[group], visible[1]) + assert torch.equal(logits[group][valid], reference), (rows, length, group) + + lengths.fill_(length) + positions.fill_(length - 1) + timings = {group: [] for group in groups} + for cycle in range(5): + order = groups if cycle % 2 == 0 else list(reversed(groups)) + for group in order: + graphs[group].replay() + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + for _ in range(repeats): + graphs[group].replay() + end.record() + end.synchronize() + timings[group].append(start.elapsed_time(end) * 1000 / repeats) + return { + "rows": rows, + "length": length, + "capacity_columns": columns, + "index_heads": heads, + "head_dim": dim, + "strided_stages": strided_stages, + "exact_changed_replays": 4, + "groups": { + str(group): { + "ctas": rows + * ( + min(-group, triton.cdiv(columns, block_n)) + if group < 0 + else triton.cdiv(columns, block_n * group) + ), + "live_ctas": rows + * ( + min(-group, triton.cdiv(length // ratio, block_n)) + if group < 0 + else triton.cdiv(length // ratio, block_n * group) + ), + "median_us": statistics.median(values), + "samples_us": values, + "compiled_resources": resources[group], + } + for group, values in timings.items() + }, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 4, 8, 16]) + parser.add_argument( + "--lengths", type=int, nargs="+", default=[8192, 131072, 262144] + ) + parser.add_argument("--groups", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument("--strided-grids", type=int, nargs="*", default=[]) + parser.add_argument("--strided-stages", type=int, choices=[1, 2], default=2) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.output.exists(): + parser.error("Refusing to overwrite a previous benchmark") + if ( + 1 not in args.groups + or min(args.groups + args.rows + args.strided_grids + [args.repeats]) < 1 + ): + parser.error("Positive rows/groups/repeats and baseline group 1 required") + if not all(1 <= length <= 262144 for length in args.lengths): + parser.error("Lengths must fit the fixed 256K capacity") + assert torch.cuda.get_device_capability() == (7, 0) + results = [] + for rows in args.rows: + for length in args.lengths: + modes = args.groups + [-grid for grid in args.strided_grids] + result = run_case(rows, length, modes, args.repeats, args.strided_stages) + results.append(result) + print(json.dumps(result), flush=True) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps({"results": results}, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/flashinfer_sm70_gdn_conv.py b/benchmarks/kernels/flashinfer_sm70_gdn_conv.py new file mode 100644 index 0000000000..92b8d0775c --- /dev/null +++ b/benchmarks/kernels/flashinfer_sm70_gdn_conv.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer-derived FP16 gate projection + conv + FP32 delta-rule prototype. + +No production dispatch. A workspace is single-stream, state indices must be +unique live pool entries or negative padding. Geometry is specialized at JIT +time; no checkpoint name, max-seqs, KV dtype or TP degree is used as a gate. +""" + +import os +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[2] +SOURCE_SHA = "6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f" + + +def namespace( + hidden, q_heads, v_heads, rows_per_warp=8, two_phase=False, shared_parameters=False +): + suffix = "" if rows_per_warp == 8 else f"_r{rows_per_warp}" + suffix += "_p2" if two_phase else "" + suffix += "_shared" if shared_parameters else "" + return f"_C_flashinfer_gdn_sm70_h{hidden}_q{q_heads}_v{v_heads}{suffix}" + + +def build( + hidden=2560, + q_heads=4, + v_heads=12, + rows_per_warp=8, + two_phase=False, + shared_parameters=False, +): + from torch.utils.cpp_extension import load + + if min(hidden, q_heads, v_heads) <= 0 or v_heads % q_heads: + raise ValueError("Require positive geometry and integral GQA grouping") + if rows_per_warp not in (4, 8): + raise ValueError("Screened row tiles are 4 or 8") + if os.environ.get("TORCH_CUDA_ARCH_LIST") != "7.0": + raise RuntimeError("Set TORCH_CUDA_ARCH_LIST=7.0") + op_namespace = namespace( + hidden, q_heads, v_heads, rows_per_warp, two_phase, shared_parameters + ) + return load( + name=op_namespace[3:], + sources=[str(ROOT / "benchmarks/csrc/sm70_flashinfer_gdn_conv.cu")], + extra_include_paths=[str(ROOT / "flashinfer-sm70/include")], + extra_cuda_cflags=[ + "-O3", + "-lineinfo", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + f"-DFI_GDN_TORCH_NAMESPACE={op_namespace}", + f"-DFI_GDN_HIDDEN={hidden}", + f"-DFI_GDN_N_BA={2 * v_heads}", + f"-DFI_GDN_QKV_DIM={(2 * q_heads + v_heads) * 128}", + f"-DFI_GDN_H_Q={q_heads}", + f"-DFI_GDN_HV={v_heads}", + "-DFI_GDN_D=128", + "-DFI_GDN_CONV_WIDTH=4", + "-DFI_GDN_CONV_STATE_LEN=3", + f"-DFI_GDN_TWO_PHASE={int(two_phase)}", + f"-DFI_GDN_SHARED_PARAMETERS={int(shared_parameters)}", + ] + + ([f"-DFI_GDN_ROWS_PER_WARP={rows_per_warp}"] if rows_per_warp != 8 else []), + is_python_module=False, + verbose=True, + ) + + +class FusedGDN: + def __init__( + self, + rows, + q_heads=4, + v_heads=12, + device="cuda", + hidden=2560, + rows_per_warp=8, + two_phase=False, + shared_parameters=False, + ): + if ( + not 1 <= rows <= 64 + or min(hidden, q_heads, v_heads) <= 0 + or v_heads % q_heads + ): + raise ValueError("Require 1..64 rows, positive geometry and integral GQA") + self.run = getattr( + torch.ops, + namespace( + hidden, q_heads, v_heads, rows_per_warp, two_phase, shared_parameters + ), + ).run + self.output = torch.empty( + rows, v_heads, 128, device=device, dtype=torch.float16 + ) + self.conv_out = torch.empty( + rows, (2 * q_heads + v_heads) * 128, device=device, dtype=torch.float16 + ) + self.partial = torch.empty( + rows * 2 * v_heads * 160, device=device, dtype=torch.float32 + ) + + def __call__( + self, + hidden, + weights, + qkv, + conv_w, + conv_bias, + conv, + A_log, + dt_bias, + state, + indices, + ): + self.run( + hidden, + weights, + qkv, + conv_w, + conv_bias, + conv, + A_log, + dt_bias, + state, + indices, + self.output, + self.conv_out, + self.partial, + ) + return self.output + + +if __name__ == "__main__": + print(build()) diff --git a/benchmarks/kernels/flashinfer_sm70_hc_norm.py b/benchmarks/kernels/flashinfer_sm70_hc_norm.py new file mode 100644 index 0000000000..5c120df327 --- /dev/null +++ b/benchmarks/kernels/flashinfer_sm70_hc_norm.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer-derived fused HC combine/Gemma-norm component, not dispatch.""" + +import os +import subprocess +from pathlib import Path + +import torch + +from benchmarks.kernels.flashinfer_sm70_gdn_conv import ROOT, SOURCE_SHA + + +def build(): + from torch.utils.cpp_extension import load + + source = Path(os.environ["FLASHINFER_SOURCE"]) + if ( + subprocess.check_output( + ["git", "-C", str(source), "rev-parse", "HEAD"], text=True + ).strip() + != SOURCE_SHA + ): + raise RuntimeError("Unrecognized FlashInfer revision") + if os.environ.get("TORCH_CUDA_ARCH_LIST") != "7.0": + raise RuntimeError("Set TORCH_CUDA_ARCH_LIST=7.0") + return load( + name="flashinfer_hc_norm_sm70_v2", + sources=[str(ROOT / "benchmarks/csrc/sm70_flashinfer_hc_norm.cu")], + extra_include_paths=[ + str(ROOT / "flashinfer-sm70/include"), + str(source / "include"), + str(source / "3rdparty/cccl/thrust"), + str(source / "3rdparty/cccl/cub"), + str(source / "3rdparty/cccl/libcudacxx/include"), + ], + extra_cuda_cflags=[ + "-O3", + "-lineinfo", + "--expt-relaxed-constexpr", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + ], + is_python_module=False, + verbose=True, + ) + + +class HCNorm: + def __init__(self, residual, warps=4, registers=False): + self.combined = torch.empty_like(residual) + self.normalized = torch.empty_like(residual) + self.warps = warps + self.registers = registers + + def __call__(self, residual, block, injection, weight, eps=1e-6): + torch.ops._C_flashinfer_hc_sm70.run( + residual, + block, + injection, + weight, + self.combined, + self.normalized, + eps, + self.warps, + self.registers, + ) + return self.combined, self.normalized + + +if __name__ == "__main__": + print(build()) diff --git a/benchmarks/kernels/flashinfer_sm70_mqa.py b/benchmarks/kernels/flashinfer_sm70_mqa.py new file mode 100644 index 0000000000..4fff01db31 --- /dev/null +++ b/benchmarks/kernels/flashinfer_sm70_mqa.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build and screen the native FlashInfer-style SM70 MQA work scheduler. + +No runtime dispatch/default change. The production-source extension is built +before capture; all mutable state belongs to the caller, including the schedule. +""" + +import os +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[2] + + +def build(): + prebuilt = os.environ.get("SM70_FLASHINFER_MQA_TEST_LIBRARY") + if prebuilt: + path = str(Path(prebuilt).resolve(strict=True)) + if not hasattr(torch.ops._C_flashinfer_mqa_sm70, "run"): + torch.ops.load_library(path) + return path + from torch.utils.cpp_extension import load + + if os.environ.get("TORCH_CUDA_ARCH_LIST") != "7.0": + raise RuntimeError("Explicit TORCH_CUDA_ARCH_LIST=7.0 required") + return load( + name="flashinfer_mqa_sm70_v1", + sources=[str(ROOT / "csrc/flashinfer_sm70/qsa_mqa.cu")], + extra_include_paths=[str(ROOT / "flashinfer-sm70/include")], + extra_cuda_cflags=["-O3", "-lineinfo", "--ptxas-options=-v"], + is_python_module=False, + verbose=True, + ) + + +class FlashInferMQA: + def __init__(self, q, columns, workers): + self.workers = workers + self.logits = torch.empty( + (q.shape[0], columns), device=q.device, dtype=torch.float32 + ) + self.visible = torch.empty(q.shape[0], device=q.device, dtype=torch.int32) + self.schedule = torch.empty( + (workers + 1, 2), device=q.device, dtype=torch.int32 + ) + + def __call__(self, q, k, table, requests, positions, lengths, ratio=4): + torch.ops._C_flashinfer_mqa_sm70.run( + q, + k, + table, + requests, + positions, + lengths, + self.logits, + self.visible, + self.schedule, + ratio, + q.shape[2] ** 0.5, + self.workers, + ) + return self.logits, self.visible + + +if __name__ == "__main__": + print(build()) diff --git a/benchmarks/kernels/flashinfer_sm70_qsa.py b/benchmarks/kernels/flashinfer_sm70_qsa.py new file mode 100644 index 0000000000..30b2311466 --- /dev/null +++ b/benchmarks/kernels/flashinfer_sm70_qsa.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark-only instantiation of pinned FlashInfer CUDA QSA decode. + +There is no model dispatch change. All mutable buffers belong to each instance; +construct/warm before CUDA Graph capture. The upstream Python SM75+ gate is not +disabled, and the previous Flash-V100 or Triton attention is never called here. +""" + +import os +import subprocess +from pathlib import Path + +import torch + +UPSTREAM_SHA = "6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f" +CCCL_SHA = "16bd510c9b712e82b0ab6cbb630d8e29ba1f7116" +ROOT = Path(__file__).resolve().parents[2] + + +def build(*, compatibility=False): + from torch.utils.cpp_extension import load + + source = Path( + os.environ.get( + "FLASHINFER_SM70_QSA_SOURCE", ROOT / ".deps/flashinfer-6c14bbd5ff34" + ) + ).resolve() + sha = subprocess.check_output( + ["git", "-C", str(source), "rev-parse", "HEAD"], text=True + ).strip() + if sha != UPSTREAM_SHA: + raise RuntimeError(f"Expected FlashInfer {UPSTREAM_SHA}, got {sha}") + if subprocess.check_output( + ["git", "-C", str(source), "diff", "HEAD", "--", "include"], text=True + ).strip(): + raise RuntimeError("Upstream CUDA headers must be unmodified") + cccl = source / "3rdparty/cccl" + if ( + subprocess.check_output( + ["git", "-C", str(cccl), "rev-parse", "HEAD"], text=True + ).strip() + != CCCL_SHA + or subprocess.check_output( + ["git", "-C", str(cccl), "diff", "HEAD"], text=True + ).strip() + ): + raise RuntimeError("Expected unmodified pinned CCCL submodule") + if os.environ.get("TORCH_CUDA_ARCH_LIST") != "7.0": + raise RuntimeError("Set TORCH_CUDA_ARCH_LIST=7.0 explicitly") + return load( + name="flashinfer_qsa_sm70_compat_v1" + if compatibility + else "flashinfer_qsa_sm70_v1", + sources=[str(ROOT / "benchmarks/csrc/sm70_flashinfer_qsa_decode.cu")], + extra_include_paths=[ + str(ROOT / "flashinfer-sm70/include"), + str(source / "include"), + str(source / "3rdparty/cccl/cub"), + str(source / "3rdparty/cccl/thrust"), + str(source / "3rdparty/cccl/libcudacxx/include"), + ], + extra_cuda_cflags=[ + "-O3", + "-lineinfo", + "--expt-relaxed-constexpr", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + f"-DFI_QSA_WMMA_COMPAT={int(compatibility)}", + ], + is_python_module=False, + verbose=True, + ) + + +class FlashInferQSA: + """Persistent single-stream workspace; use distinct instances per stream.""" + + def __init__(self, q, selection_width, splits, *, compatibility=False): + if q.ndim != 3 or q.shape[2] != 256 or q.dtype != torch.float16: + raise ValueError("QSA prototype needs FP16 [rows, heads, 256] queries") + if not 1 <= splits <= 64 or selection_width <= 0: + raise ValueError("Require positive selection width and 1..64 splits") + rows, heads, dim = q.shape + if rows <= 0 or not 1 <= heads <= 32: + raise ValueError("Require positive rows and 1..32 heads") + width = ((selection_width + splits - 1) // splits) * splits + self.splits = splits + self.offsets = torch.empty((rows, width), device=q.device, dtype=torch.int64) + self.metadata = torch.empty( + rows + 2 + 2 * rows * splits, device=q.device, dtype=torch.int32 + ) + self.zero = torch.zeros(256, device=q.device, dtype=torch.float16) + self.partial = torch.empty( + (rows, splits, heads, dim), device=q.device, dtype=torch.float32 + ) + self.lse = torch.empty( + (rows, splits, heads), device=q.device, dtype=torch.float32 + ) + self.output = torch.empty_like(q, memory_format=torch.contiguous_format) + namespace = ( + torch.ops._C_flashinfer_qsa_sm70_compat + if compatibility + else torch.ops._C_flashinfer_qsa_sm70 + ) + self.run = namespace.run + + def __call__(self, q, k, v, indices, table, requests): + self.run( + q, + k, + v, + indices, + table, + requests, + self.offsets, + self.metadata, + self.zero, + self.partial, + self.lse, + self.output, + self.splits, + ) + return self.output + + +if __name__ == "__main__": + print(build()) diff --git a/benchmarks/kernels/sm70_hc_push_gather.cuh b/benchmarks/kernels/sm70_hc_push_gather.cuh new file mode 100644 index 0000000000..4328c85fd4 --- /dev/null +++ b/benchmarks/kernels/sm70_hc_push_gather.cuh @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Preserve benchmark aliases while sharing the screened production kernels. +#include "../../csrc/sm70_hc_batch.cuh" + +TORCH_LIBRARY_FRAGMENT(_C_custom_ar_flashnext, m) { + m.def( + "hc_down_gather(int ptr, Tensor input, Tensor! injection, Tensor! lora) " + "-> ()"); + m.impl("hc_down_gather", torch::kCUDA, &sm70_hc_batch::run); + m.def("hc_mix_gather(int ptr, Tensor gate, Tensor x, Tensor! output) -> ()"); + m.impl("hc_mix_gather", torch::kCUDA, &sm70_hc_batch::run); +} diff --git a/benchmarks/kernels/sm70_hc_push_sidecar.cu b/benchmarks/kernels/sm70_hc_push_sidecar.cu new file mode 100644 index 0000000000..629af11fcb --- /dev/null +++ b/benchmarks/kernels/sm70_hc_push_sidecar.cu @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// Keep all communicator lifecycle and native features in one owning DSO. +// Bindings mirror the CUDA custom_ar block in csrc/torch_bindings.cpp. +#include "../../csrc/custom_all_reduce.cu" +#include "sm70_hc_push_gather.cuh" + +TORCH_LIBRARY(_C_custom_ar_flashnext, custom_ar) { + // Custom all-reduce kernels + custom_ar.def( + "init_custom_ar(int[] ipc_tensors, Tensor rank_data, " + "int rank, bool fully_connected) -> int"); + custom_ar.impl("init_custom_ar", torch::kCUDA, &init_custom_ar); + custom_ar.def( + "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ar.impl("all_reduce", torch::kCUDA, &all_reduce); + custom_ar.def( + "sm70_tp2_all_reduce_gemma_rms_norm(int fa, Tensor inp, Tensor " + "residual, Tensor weight, Tensor! normalized_out, Tensor! residual_out, " + "int reg_buffer, int reg_buffer_sz_bytes, float epsilon) -> ()"); + custom_ar.impl("sm70_tp2_all_reduce_gemma_rms_norm", torch::kCUDA, + &sm70_tp2_all_reduce_gemma_rms_norm); + custom_ar.def( + "sm70_tp4_all_reduce_gemma_rms_norm(int fa, Tensor inp, Tensor " + "residual, Tensor weight, Tensor! normalized_out, Tensor! residual_out, " + "int reg_buffer, int reg_buffer_sz_bytes, float epsilon) -> ()"); + custom_ar.impl("sm70_tp4_all_reduce_gemma_rms_norm", torch::kCUDA, + &sm70_tp4_all_reduce_gemma_rms_norm); + custom_ar.def( + "sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather(int fa, Tensor inp, " + "Tensor residual, Tensor weight, Tensor! normalized_out, Tensor! " + "residual_out, int reg_input_buffer, int reg_output_buffer, int " + "reg_buffer_sz_bytes, float epsilon) -> ()"); + custom_ar.impl("sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather", + torch::kCUDA, + &sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather); + custom_ar.def( + "all_reduce_sum2(int fa, Tensor inp_a, Tensor inp_b, Tensor! out) -> ()"); + custom_ar.impl("all_reduce_sum2", torch::kCUDA, &all_reduce_sum2); + custom_ar.def( + "sm70_qwen38_hc_batch_down(int ptr, Tensor input, Tensor! injection, " + "Tensor! lora) -> ()"); + custom_ar.impl("sm70_qwen38_hc_batch_down", torch::kCUDA, + &sm70_qwen38_hc_batch_down); + custom_ar.def( + "sm70_qwen38_hc_batch_mix(int ptr, Tensor gate, Tensor branches, " + "Tensor! output) -> ()"); + custom_ar.impl("sm70_qwen38_hc_batch_mix", torch::kCUDA, + &sm70_qwen38_hc_batch_mix); + custom_ar.def( + "sm70_qwen38_hc_down_allgather(int fa, Tensor inp, Tensor! out) -> ()"); + custom_ar.impl("sm70_qwen38_hc_down_allgather", torch::kCUDA, + &sm70_qwen38_hc_down_allgather); + custom_ar.def( + "sm70_qwen38_hc_gate_mix(int fa, Tensor local_gate, Tensor branches, " + "Tensor! out) -> ()"); + custom_ar.impl("sm70_qwen38_hc_gate_mix", torch::kCUDA, + &sm70_qwen38_hc_gate_mix); + custom_ar.def( + "sm70_qwen38_hc_output_allgather(int fa, Tensor local_block, " + "Tensor! out) -> ()"); + custom_ar.impl("sm70_qwen38_hc_output_allgather", torch::kCUDA, + &sm70_qwen38_hc_output_allgather); + custom_ar.def( + "sm70_qwen38_hc_up_mix_allgather(int fa, Tensor lora, Tensor weight, " + "Tensor branches, Tensor! out) -> ()"); + custom_ar.impl("sm70_qwen38_hc_up_mix_allgather", torch::kCUDA, + &sm70_qwen38_hc_up_mix_allgather); + custom_ar.def( + "top1_argmax(int fa, Tensor input_pair, Tensor! output, int reg_buffer, " + "int reg_buffer_sz_bytes) -> ()"); + custom_ar.impl("top1_argmax", torch::kCUDA, &top1_argmax); + custom_ar.def( + "tile_runtime_all_reduce(int fa, Tensor inp, Tensor! out, int " + "reg_buffer, int reg_buffer_sz_bytes, int tile_numel, int " + "engine_blocks, int compute_iters) -> ()"); + custom_ar.impl("tile_runtime_all_reduce", torch::kCUDA, + &tile_runtime_all_reduce); + custom_ar.def( + "tile_runtime_all_reduce_engine(int fa, Tensor inp, Tensor! out, int " + "reg_buffer, int reg_buffer_sz_bytes, int tile_numel, int " + "producer_blocks, int reducer_blocks, int compute_iters) -> ()"); + custom_ar.impl("tile_runtime_all_reduce_engine", torch::kCUDA, + &tile_runtime_all_reduce_engine); + custom_ar.def( + "tile_runtime_wait_reduce(int fa, Tensor staging, Tensor! out, " + "int tile_numel, int reducer_blocks) -> ()"); + custom_ar.impl("tile_runtime_wait_reduce", torch::kCUDA, + &tile_runtime_wait_reduce); + + custom_ar.def("dispose", &dispose); + custom_ar.def("meta_size", &meta_size); + custom_ar.def("sm70_tp4_push_allreduce_buffer_size", + &sm70_tp4_push_allreduce_buffer_size); + custom_ar.def("sm70_tp8_hierarchical_push_allreduce_buffer_size", + &sm70_tp8_hierarchical_push_allreduce_buffer_size); + + custom_ar.def("register_buffer", ®ister_buffer); + custom_ar.def("register_sm70_tp4_push_allreduce_buffer", + ®ister_sm70_tp4_push_allreduce_buffer); + custom_ar.def("register_sm70_tp8_hierarchical_push_allreduce_buffer", + ®ister_sm70_tp8_hierarchical_push_allreduce_buffer); + custom_ar.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta); + custom_ar.def("register_graph_buffers", ®ister_graph_buffers); + + custom_ar.def("allocate_shared_buffer_and_handle", + &allocate_shared_buffer_and_handle); + custom_ar.def("open_mem_handle(Tensor mem_handle) -> int", &open_mem_handle); + custom_ar.impl("open_mem_handle", torch::kCPU, &open_mem_handle); + + custom_ar.def("free_shared_buffer", &free_shared_buffer); +} diff --git a/benchmarks/kernels/sm70_moe_scale_layout.cu b/benchmarks/kernels/sm70_moe_scale_layout.cu new file mode 100644 index 0000000000..cc120a858a --- /dev/null +++ b/benchmarks/kernels/sm70_moe_scale_layout.cu @@ -0,0 +1,307 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Benchmark-only scale tile-major locality; arithmetic/weights unchanged. +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int kMaxRoutes = 160; +constexpr int kPack = 8; +constexpr int kExperts = 512; +constexpr int kChunks = kMaxRoutes / kPack; + +// Integer atomics only: route order within a pack is immaterial because both +// projections scatter back to the original route before the unchanged W2. +__global__ void plan_kernel(const int32_t* ids, int32_t* rows, int32_t* experts, + int32_t* sizes, int32_t* total, int routes) { + __shared__ int counts[kExperts + 1]; + __shared__ int groups[kExperts + 1][kChunks]; + const int t = threadIdx.x; + for (int e = t; e <= kExperts; e += blockDim.x) counts[e] = 0; + if (t == 0) *total = 0; + __syncthreads(); + int expert = 0, ordinal = 0; + if (t < routes) { + expert = ids[t]; + if (expert < 0 || expert >= kExperts) expert = kExperts; + ordinal = atomicAdd(counts + expert, 1); + } + __syncthreads(); + if (t < routes && ordinal % kPack == 0) { + const int group = atomicAdd(total, 1); + groups[expert][ordinal / kPack] = group; + experts[group] = expert; + sizes[group] = min(kPack, counts[expert] - ordinal); + } + __syncthreads(); + if (t < routes) { + const int group = groups[expert][ordinal / kPack]; + rows[group * kPack + ordinal % kPack] = t; + } +} + +__device__ __forceinline__ void decode(unsigned packed, half2 scale, + half2* out) { + constexpr unsigned sign = 0x80008000u, em = 0x0e000e00u; + unsigned v[4] = {((packed << 12) & sign) | ((packed << 9) & em), + ((packed << 8) & sign) | ((packed << 5) & em), + ((packed << 4) & sign) | ((packed << 1) & em), + (packed & sign) | ((packed >> 3) & em)}; +#pragma unroll + for (int i = 0; i < 4; ++i) + out[i] = __hmul2(*reinterpret_cast(v + i), scale); +} + +#define PACKED_MMA(C, A0, A1, B0, B1) \ + asm volatile( \ + "mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " \ + "{%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, " \ + "{%0,%1,%2,%3,%4,%5,%6,%7};\n" \ + : "+f"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3]), "+f"(C[4]), \ + "+f"(C[5]), "+f"(C[6]), "+f"(C[7]) \ + : "r"(A0), "r"(A1), "r"(B0), "r"(B1)) + +template +__global__ void w13_kernel(const half* x, const uint32_t* weights, + const half* scales, const int32_t* rows, + const int32_t* experts, const int32_t* sizes, + const int32_t* total, half* out) { + // Split within the CTA: no floating-point atomics or global partial tensor. + __shared__ float partial[2][Split][kPack][32]; + __shared__ half projected[2][kPack][32]; + const int group_id = blockIdx.y; + if (group_id >= *total) return; + const int count = sizes[group_id], expert = experts[group_id]; + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + const int projection = warp / Split, split = warp % Split; + const int tile = + Interleaved ? blockIdx.x * 2 + projection : blockIdx.x + projection * 5; + const int mma_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int quad = (lane >> 2) & 3; + const int col = quad * 8 + mma_row; + const int route = mma_row < count ? rows[group_id * kPack + mma_row] : 0; + float accum[8] = {}; + if (expert < kExperts) { + const uint32_t* w = weights + static_cast(expert) * 2560 * 40; + const half* s = scales + static_cast(expert) * 160 * 320; + const half* input = x + static_cast(route / 10) * 2560; +#pragma unroll 4 + for (int g = split * (160 / Split); g < (split + 1) * (160 / Split); ++g) { + const size_t offset = + (static_cast(tile) * 320 + g * 2) * 32 + col; + const half scalar = __hmul(__ldg(s + (tile * 160 + g) * 32 + col), + __float2half_rn(16384.0f)); + const half2 scale = __halves2half2(scalar, scalar); + half2 decoded[8]; + decode(__ldcs(w + offset), scale, decoded); + decode(__ldcs(w + offset + 32), scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + uint4 lo = make_uint4(0, 0, 0, 0), hi = make_uint4(0, 0, 0, 0); + if (mma_row < count) { + lo = *reinterpret_cast(input + g * 16); + hi = *reinterpret_cast(input + g * 16 + 8); + } + PACKED_MMA(accum, lo.x, lo.y, b[0], b[1]); + PACKED_MMA(accum, lo.z, lo.w, b[2], b[3]); + PACKED_MMA(accum, hi.x, hi.y, b[4], b[5]); + PACKED_MMA(accum, hi.z, hi.w, b[6], b[7]); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int r = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + partial[projection][split][r][quad * 8 + c] = accum[i]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 2 * kPack * 32; idx += blockDim.x) { + const int p = idx / (kPack * 32), r = idx / 32 % kPack, c = idx % 32; + // FP16 materialization is retained before SiLU, then again before the + // multiplication. Split>1 changes FP32 association, not quantization. + float value = 0; +#pragma unroll + for (int s = 0; s < Split; ++s) value += partial[p][s][r][c]; + projected[p][r][c] = __float2half_rn(value); + } + __syncthreads(); + for (int idx = threadIdx.x; idx < count * 32; idx += blockDim.x) { + const int r = idx / 32, c = idx % 32; + const int p = Interleaved ? c / 16 : 0; + const int pc = Interleaved ? c % 16 * 2 : c; + const half gate = projected[p][r][pc]; + const half up = Interleaved ? projected[p][r][pc + 1] : projected[1][r][c]; + const float gf = __half2float(gate); + const half activated = __float2half_rn(gf / (1.0f + expf(-gf))); + out[static_cast(rows[group_id * kPack + r]) * 160 + + blockIdx.x * 32 + c] = __hmul(activated, up); + } +} + +void run(torch::Tensor out, torch::Tensor x, torch::Tensor w, torch::Tensor s, + torch::Tensor ids, torch::Tensor rows, torch::Tensor experts, + torch::Tensor sizes, torch::Tensor total, int64_t split, + bool interleaved) { + const c10::cuda::CUDAGuard guard(x.device()); + const int routes = x.size(0) * 10; + TORCH_CHECK(x.dim() == 2 && x.size(1) == 2560 && routes > 0 && routes <= 160); + for (const auto& t : {out, x, w, s, ids, rows, experts, sizes, total}) { + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous()); + } + TORCH_CHECK(x.scalar_type() == at::kHalf && s.scalar_type() == at::kHalf && + out.scalar_type() == at::kHalf && w.scalar_type() == at::kInt); + for (const auto& t : {ids, rows, experts, sizes, total}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(ids.numel() == routes && rows.numel() >= routes * kPack && + experts.numel() >= routes && sizes.numel() >= routes && + total.numel() == 1 && out.numel() == routes * 160 && + w.numel() == 512 * 2560 * 40 && s.numel() == 512 * 160 * 320); + const auto stream = at::cuda::getCurrentCUDAStream(x.get_device()); + plan_kernel<<<1, 256, 0, stream>>>( + ids.data_ptr(), rows.data_ptr(), + experts.data_ptr(), sizes.data_ptr(), + total.data_ptr(), routes); +#define LAUNCH(S, I) \ + w13_kernel<<>>( \ + reinterpret_cast(x.data_ptr()), \ + reinterpret_cast(w.data_ptr()), \ + reinterpret_cast(s.data_ptr()), rows.data_ptr(), \ + experts.data_ptr(), sizes.data_ptr(), \ + total.data_ptr(), reinterpret_cast(out.data_ptr())) +#define CASE(S) \ + case S: \ + if (interleaved) { \ + LAUNCH(S, true); \ + } else { \ + LAUNCH(S, false); \ + } \ + break + switch (split) { + CASE(1); + CASE(2); + CASE(4); + CASE(5); + CASE(8); + default: + TORCH_CHECK(false, "Unsupported split"); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +#undef CASE +#undef LAUNCH +} + +// All groups, including singletons, use one kernel. The earlier prototype +// launched separate repeated/singleton kernels; their overhead erased reuse. +__global__ void w2_kernel(const half* x, const uint32_t* weights, + const half* scales, const int32_t* rows, + const int32_t* experts, const int32_t* sizes, + const int32_t* total, half* out) { + const int group = blockIdx.y * 4 + threadIdx.x / 32; + if (group >= *total) return; + const int count = sizes[group], expert = experts[group]; + const int lane = threadIdx.x % 32, quad = (lane >> 2) & 3; + const int r = (lane & 3) + ((lane & 16) ? 4 : 0), col = quad * 8 + r; + const int route = r < count ? rows[group * kPack + r] : 0; + float accum[8] = {}; + if (expert < kExperts) { + const uint32_t* w = weights + static_cast(expert) * 160 * 320; + const half* s = scales + static_cast(expert) * 10 * 2560; + const half* input = x + static_cast(route) * 160; +#pragma unroll + for (int g = 0; g < 10; ++g) { + const int offset = (blockIdx.x * 20 + g * 2) * 32 + col; + const half scalar = __hmul(__ldg(s + (blockIdx.x * 10 + g) * 32 + col), + __float2half_rn(16384.0f)); + const half2 scale = __halves2half2(scalar, scalar); + half2 decoded[8]; + decode(__ldcs(w + offset), scale, decoded); + decode(__ldcs(w + offset + 32), scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + uint4 lo = make_uint4(0, 0, 0, 0), hi = make_uint4(0, 0, 0, 0); + if (r < count) { + lo = *reinterpret_cast(input + g * 16); + hi = *reinterpret_cast(input + g * 16 + 8); + } + PACKED_MMA(accum, lo.x, lo.y, b[0], b[1]); + PACKED_MMA(accum, lo.z, lo.w, b[2], b[3]); + PACKED_MMA(accum, hi.x, hi.y, b[4], b[5]); + PACKED_MMA(accum, hi.z, hi.w, b[6], b[7]); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int row = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + if (row < count) { + const int dst = rows[group * kPack + row]; + out[static_cast(dst) * 2560 + blockIdx.x * 32 + quad * 8 + c] = + __float2half_rn(accum[i]); + } + } +} + +__global__ void reduce_kernel(const half* routed, const float* weights, + half* out) { + const int token = blockIdx.y, col = blockIdx.x * 256 + threadIdx.x; + float result = 0; +#pragma unroll + for (int slot = 0; slot < 10; ++slot) + result = fmaf(__half2float(routed[(token * 10 + slot) * 2560 + col]), + weights[token * 10 + slot], result); + out[token * 2560 + col] = __float2half_rn(result); +} + +void w2(torch::Tensor out, torch::Tensor routed, torch::Tensor x, + torch::Tensor w, torch::Tensor s, torch::Tensor topk, + torch::Tensor rows, torch::Tensor experts, torch::Tensor sizes, + torch::Tensor total) { + const c10::cuda::CUDAGuard guard(x.device()); + TORCH_CHECK(x.dim() == 2 && x.size(1) == 160 && x.size(0) % 10 == 0); + const int routes = x.size(0), tokens = routes / 10; + TORCH_CHECK(tokens >= 1 && tokens <= 16); + for (const auto& t : + {out, routed, x, w, s, topk, rows, experts, sizes, total}) + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous()); + for (const auto& t : {out, routed, x, s}) + TORCH_CHECK(t.scalar_type() == at::kHalf); + for (const auto& t : {w, rows, experts, sizes, total}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(topk.scalar_type() == at::kFloat && topk.numel() == routes && + out.numel() == tokens * 2560 && routed.numel() == routes * 2560 && + rows.numel() >= routes * 8 && experts.numel() >= routes && + sizes.numel() >= routes && total.numel() == 1 && + w.numel() == 512 * 160 * 320 && s.numel() == 512 * 10 * 2560); + const auto stream = at::cuda::getCurrentCUDAStream(x.get_device()); + w2_kernel<<>>( + reinterpret_cast(x.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(s.data_ptr()), rows.data_ptr(), + experts.data_ptr(), sizes.data_ptr(), + total.data_ptr(), reinterpret_cast(routed.data_ptr())); + reduce_kernel<<>>( + reinterpret_cast(routed.data_ptr()), topk.data_ptr(), + reinterpret_cast(out.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace + +TORCH_LIBRARY_FRAGMENT(_C_moe_scale_layout, m) { + m.def( + "nvfp4_grouped_w13_sm70_out(Tensor(a!) out, Tensor x, Tensor w, Tensor " + "s, Tensor ids, " + "Tensor(b!) rows, Tensor(c!) experts, Tensor(d!) sizes, Tensor(e!) " + "total, " + "int split, bool interleaved) -> ()"); + m.def( + "nvfp4_grouped_w2_sm70_out(Tensor(a!) out, Tensor(b!) routed, Tensor x, " + "Tensor w, Tensor s, " + "Tensor topk, Tensor rows, Tensor experts, Tensor sizes, Tensor total) " + "-> ()"); +} +TORCH_LIBRARY_IMPL(_C_moe_scale_layout, CUDA, m) { + m.impl("nvfp4_grouped_w13_sm70_out", &run); + m.impl("nvfp4_grouped_w2_sm70_out", &w2); +} diff --git a/benchmarks/kernels/sm70_moe_w13_paired.cu b/benchmarks/kernels/sm70_moe_w13_paired.cu new file mode 100644 index 0000000000..89679cce41 --- /dev/null +++ b/benchmarks/kernels/sm70_moe_w13_paired.cu @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Benchmark only. A warp reuses one A fragment for gate and up projections. +// Preserve both independent accumulation sequences and the ordered Split-K sum. +#define _C _C_moe_pair_reference +#include "../../csrc/sm70_turbomind/ops/nvfp4_grouped_decode_sm70.cu" +#undef _C +namespace { +template +__global__ void paired_w13_kernel(const half* x, const uint32_t* weights, + const half* scales, const int32_t* rows, + const int32_t* experts, const int32_t* sizes, + const int32_t* total, half* out) { + // Split within the CTA: no floating-point atomics or global partial tensor. + __shared__ float partial[2][Split][kPack][32]; + __shared__ half projected[2][kPack][32]; + const int group_id = blockIdx.y; + if (group_id >= *total) return; + const int count = sizes[group_id], expert = experts[group_id]; + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + const int split = warp; + const int mma_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int quad = (lane >> 2) & 3; + const int col = quad * 8 + mma_row; + const int route = mma_row < count ? rows[group_id * kPack + mma_row] : 0; + float accum[2][8] = {}; + if (expert < kExperts) { + const uint32_t* w = weights + static_cast(expert) * 2560 * 40; + const half* s = scales + static_cast(expert) * 160 * 320; + const half* input = x + static_cast(route / 10) * 2560; +#pragma unroll 4 + for (int g = split * (160 / Split); g < (split + 1) * (160 / Split); ++g) { + uint4 lo = make_uint4(0, 0, 0, 0), hi = make_uint4(0, 0, 0, 0); + if (mma_row < count) { + lo = *reinterpret_cast(input + g * 16); + hi = *reinterpret_cast(input + g * 16 + 8); + } + +#pragma unroll + for (int projection = 0; projection < 2; ++projection) { + const int tile = Interleaved ? blockIdx.x * 2 + projection + : blockIdx.x + projection * 5; + + const size_t offset = + (static_cast(tile) * 320 + g * 2) * 32 + col; + const half scalar = __hmul(__ldg(s + (g * 10 + tile) * 32 + col), + __float2half_rn(16384.0f)); + const half2 scale = __halves2half2(scalar, scalar); + half2 decoded[8]; + decode(__ldcs(w + offset), scale, decoded); + decode(__ldcs(w + offset + 32), scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + PACKED_MMA(accum[projection], lo.x, lo.y, b[0], b[1]); + PACKED_MMA(accum[projection], lo.z, lo.w, b[2], b[3]); + PACKED_MMA(accum[projection], hi.x, hi.y, b[4], b[5]); + PACKED_MMA(accum[projection], hi.z, hi.w, b[6], b[7]); + } + } + } +#pragma unroll + for (int projection = 0; projection < 2; ++projection) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int r = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + partial[projection][split][r][quad * 8 + c] = accum[projection][i]; + } + } + __syncthreads(); + for (int idx = threadIdx.x; idx < 2 * kPack * 32; idx += blockDim.x) { + const int p = idx / (kPack * 32), r = idx / 32 % kPack, c = idx % 32; + // FP16 materialization is retained before SiLU, then again before the + // multiplication. Split>1 changes FP32 association, not quantization. + float value = 0; +#pragma unroll + for (int s = 0; s < Split; ++s) value += partial[p][s][r][c]; + projected[p][r][c] = __float2half_rn(value); + } + __syncthreads(); + for (int idx = threadIdx.x; idx < count * 32; idx += blockDim.x) { + const int r = idx / 32, c = idx % 32; + const int p = Interleaved ? c / 16 : 0; + const int pc = Interleaved ? c % 16 * 2 : c; + const half gate = projected[p][r][pc]; + const half up = Interleaved ? projected[p][r][pc + 1] : projected[1][r][c]; + const float gf = __half2float(gate); + const half activated = __float2half_rn(gf / (1.0f + expf(-gf))); + out[static_cast(rows[group_id * kPack + r]) * 160 + + blockIdx.x * 32 + c] = __hmul(activated, up); + } +} + +void paired_w13(torch::Tensor out, torch::Tensor x, torch::Tensor w, + torch::Tensor s, torch::Tensor ids, torch::Tensor rows, + torch::Tensor experts, torch::Tensor sizes, torch::Tensor total, + int64_t split, bool interleaved) { + const c10::cuda::CUDAGuard guard(x.device()); + const int routes = x.size(0) * 10; + TORCH_CHECK(x.dim() == 2 && x.size(1) == 2560 && routes > 0 && routes <= 160); + for (const auto& t : {out, x, w, s, ids, rows, experts, sizes, total}) { + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous()); + } + TORCH_CHECK(x.scalar_type() == at::kHalf && s.scalar_type() == at::kHalf && + out.scalar_type() == at::kHalf && w.scalar_type() == at::kInt); + for (const auto& t : {ids, rows, experts, sizes, total}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(ids.numel() == routes && rows.numel() >= routes * kPack && + experts.numel() >= routes && sizes.numel() >= routes && + total.numel() == 1 && out.numel() == routes * 160 && + w.numel() == 512 * 2560 * 40 && s.numel() == 512 * 160 * 320); + const auto stream = at::cuda::getCurrentCUDAStream(x.get_device()); + plan_kernel<<<1, 256, 0, stream>>>( + ids.data_ptr(), rows.data_ptr(), + experts.data_ptr(), sizes.data_ptr(), + total.data_ptr(), routes); +#define LAUNCH(S, I) \ + paired_w13_kernel<<>>( \ + reinterpret_cast(x.data_ptr()), \ + reinterpret_cast(w.data_ptr()), \ + reinterpret_cast(s.data_ptr()), rows.data_ptr(), \ + experts.data_ptr(), sizes.data_ptr(), \ + total.data_ptr(), reinterpret_cast(out.data_ptr())) +#define CASE(S) \ + case S: \ + if (interleaved) { \ + LAUNCH(S, true); \ + } else { \ + LAUNCH(S, false); \ + } \ + break + switch (split) { + CASE(1); + CASE(2); + CASE(4); + CASE(5); + CASE(8); + default: + TORCH_CHECK(false, "Unsupported split"); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +#undef CASE +#undef LAUNCH +} + +template +std::vector resource_info() { + cudaFuncAttributes attr; + AT_CUDA_CHECK(cudaFuncGetAttributes(&attr, paired_w13_kernel)); + int blocks = 0; + AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks, paired_w13_kernel, 32 * S, 0)); + return {attr.numRegs, static_cast(attr.sharedSizeBytes), + static_cast(attr.localSizeBytes), blocks}; +} +std::vector resources(int64_t split) { + switch (split) { + case 4: + return resource_info<4>(); + case 5: + return resource_info<5>(); + case 8: + return resource_info<8>(); + default: + TORCH_CHECK(false, "resource query supports production screen splits"); + } +} +} // namespace +TORCH_LIBRARY_FRAGMENT(_C_moe_pair, m) { + m.def( + "run(Tensor(a!) out, Tensor x, Tensor w, Tensor s, Tensor ids, " + "Tensor(b!) rows, Tensor(c!) experts, Tensor(d!) sizes, Tensor(e!) " + "total, " + "int split, bool interleaved) -> ()"); + m.def("resources(int split) -> int[]", &resources); +} +TORCH_LIBRARY_IMPL(_C_moe_pair, CUDA, m) { m.impl("run", &paired_w13); } diff --git a/benchmarks/kernels/sm70_moe_w13_single_tile.cu b/benchmarks/kernels/sm70_moe_w13_single_tile.cu new file mode 100644 index 0000000000..e833253555 --- /dev/null +++ b/benchmarks/kernels/sm70_moe_w13_single_tile.cu @@ -0,0 +1,173 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Benchmark only: interleaved gate/up pairs fit in a single physical N32 tile. +// Split the original two-tile CTA without changing any accumulation sequence. +#define _C _C_moe_single_tile_reference +#include "../../csrc/sm70_turbomind/ops/nvfp4_grouped_decode_sm70.cu" +#undef _C +namespace { +template +__global__ +#ifdef MOE_SINGLE_TILE_BOUND +__launch_bounds__(32 * Split, Split == 8 ? 6 : 1) +#endif + void single_tile_w13_kernel(const half* x, const uint32_t* weights, + const half* scales, const int32_t* rows, + const int32_t* experts, const int32_t* sizes, + const int32_t* total, half* out) { + // Split within the CTA: no floating-point atomics or global partial tensor. + __shared__ float partial[Split][kPack][32]; + __shared__ half projected[kPack][32]; + const int group_id = blockIdx.y; + if (group_id >= *total) return; + const int count = sizes[group_id], expert = experts[group_id]; + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + const int split = warp; + const int tile = blockIdx.x; + const int mma_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int quad = (lane >> 2) & 3; + const int col = quad * 8 + mma_row; + const int route = mma_row < count ? rows[group_id * kPack + mma_row] : 0; + float accum[8] = {}; + if (expert < kExperts) { + const uint32_t* w = weights + static_cast(expert) * 2560 * 40; + const half* s = scales + static_cast(expert) * 160 * 320; + const half* input = x + static_cast(route / 10) * 2560; +#pragma unroll 4 + for (int g = split * (160 / Split); g < (split + 1) * (160 / Split); ++g) { + const size_t offset = + (static_cast(tile) * 320 + g * 2) * 32 + col; + const half scalar = __hmul(__ldg(s + (g * 10 + tile) * 32 + col), + __float2half_rn(16384.0f)); + const half2 scale = __halves2half2(scalar, scalar); + half2 decoded[8]; + decode(__ldcs(w + offset), scale, decoded); + decode(__ldcs(w + offset + 32), scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + uint4 lo = make_uint4(0, 0, 0, 0), hi = make_uint4(0, 0, 0, 0); + if (mma_row < count) { + lo = *reinterpret_cast(input + g * 16); + hi = *reinterpret_cast(input + g * 16 + 8); + } + PACKED_MMA(accum, lo.x, lo.y, b[0], b[1]); + PACKED_MMA(accum, lo.z, lo.w, b[2], b[3]); + PACKED_MMA(accum, hi.x, hi.y, b[4], b[5]); + PACKED_MMA(accum, hi.z, hi.w, b[6], b[7]); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int r = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + partial[split][r][quad * 8 + c] = accum[i]; + } + __syncthreads(); + for (int idx = threadIdx.x; idx < kPack * 32; idx += blockDim.x) { + const int r = idx / 32, c = idx % 32; + // FP16 materialization is retained before SiLU, then again before the + // multiplication. Split>1 changes FP32 association, not quantization. + float value = 0; +#pragma unroll + for (int s = 0; s < Split; ++s) value += partial[s][r][c]; + projected[r][c] = __float2half_rn(value); + } + __syncthreads(); + for (int idx = threadIdx.x; idx < count * 16; idx += blockDim.x) { + const int r = idx / 16, c = idx % 16; + const half gate = projected[r][2 * c]; + const half up = projected[r][2 * c + 1]; + const float gf = __half2float(gate); + const half activated = __float2half_rn(gf / (1.0f + expf(-gf))); + out[static_cast(rows[group_id * kPack + r]) * 160 + + blockIdx.x * 16 + c] = __hmul(activated, up); + } +} + +void single_tile_w13(torch::Tensor out, torch::Tensor x, torch::Tensor w, + torch::Tensor s, torch::Tensor ids, torch::Tensor rows, + torch::Tensor experts, torch::Tensor sizes, + torch::Tensor total, int64_t split, bool interleaved) { + TORCH_CHECK(interleaved, + "Single-tile benchmark requires interleaved gate/up"); + const c10::cuda::CUDAGuard guard(x.device()); + const int routes = x.size(0) * 10; + TORCH_CHECK(x.dim() == 2 && x.size(1) == 2560 && routes > 0 && routes <= 160); + for (const auto& t : {out, x, w, s, ids, rows, experts, sizes, total}) { + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous()); + } + TORCH_CHECK(x.scalar_type() == at::kHalf && s.scalar_type() == at::kHalf && + out.scalar_type() == at::kHalf && w.scalar_type() == at::kInt); + for (const auto& t : {ids, rows, experts, sizes, total}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(ids.numel() == routes && rows.numel() >= routes * kPack && + experts.numel() >= routes && sizes.numel() >= routes && + total.numel() == 1 && out.numel() == routes * 160 && + w.numel() == 512 * 2560 * 40 && s.numel() == 512 * 160 * 320); + const auto stream = at::cuda::getCurrentCUDAStream(x.get_device()); + plan_kernel<<<1, 256, 0, stream>>>( + ids.data_ptr(), rows.data_ptr(), + experts.data_ptr(), sizes.data_ptr(), + total.data_ptr(), routes); +#define LAUNCH(S, I) \ + single_tile_w13_kernel<<>>( \ + reinterpret_cast(x.data_ptr()), \ + reinterpret_cast(w.data_ptr()), \ + reinterpret_cast(s.data_ptr()), rows.data_ptr(), \ + experts.data_ptr(), sizes.data_ptr(), \ + total.data_ptr(), reinterpret_cast(out.data_ptr())) +#define CASE(S) \ + case S: \ + if (interleaved) { \ + LAUNCH(S, true); \ + } else { \ + LAUNCH(S, false); \ + } \ + break + switch (split) { + CASE(1); + CASE(2); + CASE(4); + CASE(5); + CASE(8); + default: + TORCH_CHECK(false, "Unsupported split"); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +#undef CASE +#undef LAUNCH +} + +template +std::vector resource_info() { + cudaFuncAttributes attr; + AT_CUDA_CHECK(cudaFuncGetAttributes(&attr, single_tile_w13_kernel)); + int blocks = 0; + AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks, single_tile_w13_kernel, 32 * S, 0)); + return {attr.numRegs, static_cast(attr.sharedSizeBytes), + static_cast(attr.localSizeBytes), blocks}; +} +std::vector resources(int64_t split) { + switch (split) { + case 4: + return resource_info<4>(); + case 5: + return resource_info<5>(); + case 8: + return resource_info<8>(); + default: + TORCH_CHECK(false, "resource query supports production screen splits"); + } +} +} // namespace +TORCH_LIBRARY_FRAGMENT(_C_moe_single_tile, m) { + m.def( + "run(Tensor(a!) out, Tensor x, Tensor w, Tensor s, Tensor ids, " + "Tensor(b!) rows, Tensor(c!) experts, Tensor(d!) sizes, Tensor(e!) " + "total, " + "int split, bool interleaved) -> ()"); + m.def("resources(int split) -> int[]", &resources); +} +TORCH_LIBRARY_IMPL(_C_moe_single_tile, CUDA, m) { + m.impl("run", &single_tile_w13); +} diff --git a/benchmarks/kernels/sm70_moe_w2_locality.cu b/benchmarks/kernels/sm70_moe_w2_locality.cu new file mode 100644 index 0000000000..57d197f3cc --- /dev/null +++ b/benchmarks/kernels/sm70_moe_w2_locality.cu @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Benchmark-only: colocate neighbouring N tiles of one expert within a CTA. +// Reuse the production plan/reduction and an isolated reference namespace. +#define _C _C_moe_locality_reference +#include "../../csrc/sm70_turbomind/ops/nvfp4_grouped_decode_sm70.cu" +#undef _C + +namespace { +template +__global__ void locality_w2_kernel(const half* x, const uint32_t* weights, + const half* scales, const int32_t* rows, + const int32_t* experts, const int32_t* sizes, + const int32_t* total, half* out) { + const int warp = threadIdx.x / 32; + const int group = blockIdx.y * (4 / WarpTiles) + warp / WarpTiles; + const int tile = blockIdx.x * WarpTiles + warp % WarpTiles; + if (group >= *total) return; + const int count = sizes[group], expert = experts[group]; + const int lane = threadIdx.x % 32, quad = (lane >> 2) & 3; + const int r = (lane & 3) + ((lane & 16) ? 4 : 0), col = quad * 8 + r; + const int route = r < count ? rows[group * kPack + r] : 0; + float accum[8] = {}; + if (expert < kExperts) { + const uint32_t* w = weights + static_cast(expert) * 160 * 320; + const half* s = scales + static_cast(expert) * 10 * 2560; + const half* input = x + static_cast(route) * 160; +#pragma unroll + for (int g = 0; g < 10; ++g) { + const int offset = (tile * 20 + g * 2) * 32 + col; + const half scalar = __hmul(__ldg(s + (g * 80 + tile) * 32 + col), + __float2half_rn(16384.0f)); + const half2 scale = __halves2half2(scalar, scalar); + half2 decoded[8]; + decode(__ldcs(w + offset), scale, decoded); + decode(__ldcs(w + offset + 32), scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + uint4 lo = make_uint4(0, 0, 0, 0), hi = make_uint4(0, 0, 0, 0); + if (r < count) { + lo = *reinterpret_cast(input + g * 16); + hi = *reinterpret_cast(input + g * 16 + 8); + } + PACKED_MMA(accum, lo.x, lo.y, b[0], b[1]); + PACKED_MMA(accum, lo.z, lo.w, b[2], b[3]); + PACKED_MMA(accum, hi.x, hi.y, b[4], b[5]); + PACKED_MMA(accum, hi.z, hi.w, b[6], b[7]); + } + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int row = (i & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2); + if (row < count) { + const int dst = rows[group * kPack + row]; + out[static_cast(dst) * 2560 + tile * 32 + quad * 8 + c] = + __float2half_rn(accum[i]); + } + } +} + +void locality_w2(torch::Tensor out, torch::Tensor routed, torch::Tensor x, + torch::Tensor w, torch::Tensor s, torch::Tensor topk, + torch::Tensor rows, torch::Tensor experts, torch::Tensor sizes, + torch::Tensor total, int64_t warp_tiles) { + const c10::cuda::CUDAGuard guard(x.device()); + TORCH_CHECK(x.dim() == 2 && x.size(1) == 160 && x.size(0) % 10 == 0); + const int routes = x.size(0), tokens = routes / 10; + TORCH_CHECK(tokens >= 1 && tokens <= 16); + for (const auto& t : + {out, routed, x, w, s, topk, rows, experts, sizes, total}) + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.is_contiguous()); + for (const auto& t : {out, routed, x, s}) + TORCH_CHECK(t.scalar_type() == at::kHalf); + for (const auto& t : {w, rows, experts, sizes, total}) + TORCH_CHECK(t.scalar_type() == at::kInt); + TORCH_CHECK(topk.scalar_type() == at::kFloat && topk.numel() == routes && + out.numel() == tokens * 2560 && routed.numel() == routes * 2560 && + rows.numel() >= routes * 8 && experts.numel() >= routes && + sizes.numel() >= routes && total.numel() == 1 && + w.numel() == 512 * 160 * 320 && s.numel() == 512 * 10 * 2560); + const auto stream = at::cuda::getCurrentCUDAStream(x.get_device()); +#define LAUNCH(W) \ + locality_w2_kernel \ + <<>>( \ + reinterpret_cast(x.data_ptr()), \ + reinterpret_cast(w.data_ptr()), \ + reinterpret_cast(s.data_ptr()), \ + rows.data_ptr(), experts.data_ptr(), \ + sizes.data_ptr(), total.data_ptr(), \ + reinterpret_cast(routed.data_ptr())); + switch (warp_tiles) { + case 1: + LAUNCH(1); + break; + case 2: + LAUNCH(2); + break; + case 4: + LAUNCH(4); + break; + default: + TORCH_CHECK(false, "warp_tiles must be 1, 2 or 4"); + } +#undef LAUNCH + reduce_kernel<<>>( + reinterpret_cast(routed.data_ptr()), topk.data_ptr(), + reinterpret_cast(out.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +std::vector resource_info() { + cudaFuncAttributes attr; + AT_CUDA_CHECK(cudaFuncGetAttributes(&attr, locality_w2_kernel)); + int blocks = 0; + AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks, locality_w2_kernel, 128, 0)); + return {attr.numRegs, static_cast(attr.sharedSizeBytes), + static_cast(attr.localSizeBytes), blocks}; +} + +std::vector resources(int64_t warp_tiles) { + switch (warp_tiles) { + case 1: + return resource_info<1>(); + case 2: + return resource_info<2>(); + case 4: + return resource_info<4>(); + default: + TORCH_CHECK(false, "warp_tiles must be 1, 2 or 4"); + } +} +} // namespace + +TORCH_LIBRARY_FRAGMENT(_C_moe_locality, m) { + m.def( + "w2(Tensor(a!) out, Tensor(b!) routed, Tensor x, Tensor w, Tensor s, " + "Tensor topk, Tensor rows, Tensor experts, Tensor sizes, Tensor total, " + "int warp_tiles) -> ()"); + m.def("resources(int warp_tiles) -> int[]", &resources); +} +TORCH_LIBRARY_IMPL(_C_moe_locality, CUDA, m) { m.impl("w2", &locality_w2); } diff --git a/benchmarks/kernels/sm70_paired_stats.py b/benchmarks/kernels/sm70_paired_stats.py new file mode 100644 index 0000000000..c932dd7adb --- /dev/null +++ b/benchmarks/kernels/sm70_paired_stats.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paired five-block microbenchmark interval, not cross-run model evidence.""" + +import math +import statistics + + +def paired_latency_interval(control, candidate): + """Student-t interval on five interleaved log speed ratios (four DOF).""" + if len(control) != 5 or len(candidate) != 5: + raise ValueError("Exactly five paired timing blocks are required") + if any(not math.isfinite(v) or v <= 0 for v in (*control, *candidate)): + raise ValueError("Timing samples must be finite and positive") + gains = [math.log(a / b) for a, b in zip(control, candidate, strict=True)] + center = statistics.mean(gains) + radius = 2.7764451051977987 * statistics.stdev(gains) / math.sqrt(5) + return { + "method": "paired log ratio, Student-t df=4, five within-run blocks", + "geomean_reduction_pct": 100 * (1 - math.exp(-center)), + "reduction_pct_ci95": [ + 100 * (1 - math.exp(-(center - radius))), + 100 * (1 - math.exp(-(center + radius))), + ], + "positive_lower_bound": center - radius > 0, + } diff --git a/benchmarks/kernels/sm70_qsa_strided_scorer.py b/benchmarks/kernels/sm70_qsa_strided_scorer.py new file mode 100644 index 0000000000..e25115cbb2 --- /dev/null +++ b/benchmarks/kernels/sm70_qsa_strided_scorer.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark-only grid-stride scorer, derived from the NVIDIA QSA kernel. + +Only tile ownership changes; each tile retains the same dot/head reduction. +A bounded grid still visits ALL live tiles, including graph replay growth. +No production dispatcher references this experiment. +""" + +from vllm.triton_utils import tl, triton + + +@triton.jit(do_not_specialize=["num_requests"]) +def strided_qsa_mqa_paged_kernel( + q_ptr, + k_cache_ptr, + page_table_ptr, + token_to_req_ptr, + query_positions_ptr, + sequence_lengths_ptr, + visible_blocks_ptr, + logits_ptr, + stride_q_row, + stride_q_head, + stride_q_dim, + stride_cache_block, + stride_cache_token, + stride_cache_dim, + stride_table_req, + stride_table_page, + stride_logits_row, + num_rows, + num_columns, + num_pages, + num_requests, + score_divisor, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, + TILES_PER_PROG: tl.constexpr, + STAGES: tl.constexpr, + MAX_N: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +) -> None: + row = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + heads = tl.arange(0, MAX_N) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_position = tl.load(query_positions_ptr + row) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(request >= 0) & (request < num_requests), + other=0, + ) + visible = tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ) + if tl.program_id(1) == 0: + tl.store(visible_blocks_ptr + row, visible) + tile_start = tl.program_id(1) + # Top-k is bounded by visible_blocks, so columns beyond it need no value. + if tile_start * BLOCK_N >= visible: + return + tile_end = tl.cdiv(visible, BLOCK_N) + tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N)) + + # Pad the small head axis to a tensor-core-compatible N dimension. + query = tl.load( + q_ptr + + row * stride_q_row + + heads[None, :] * stride_q_head + + dims[:, None] * stride_q_dim, + mask=(heads[None, :] < NUM_HEADS) & (dims[:, None] < HEAD_DIM), + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(tile_start, tile_end, tl.num_programs(1), num_stages=STAGES): + columns = tile * BLOCK_N + column_offsets + live = columns < visible + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) + page_offset = columns % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + + safe_request * stride_table_req + + logical_page * stride_table_page, + mask=live, + other=-1, + ) + page_valid = live & (physical_page >= 0) & (physical_page < num_pages) + # physical_page * block stride can overflow int32 for large caches. + safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr + + safe_physical_page[:, None] * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :] * stride_cache_dim, + mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + scores = tl.dot(keys, query, out_dtype=tl.float32) + scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(scores, 0.0), 0.0) + score = tl.sum(scores, axis=1) / score_divisor + tl.store( + logits_ptr + row * stride_logits_row + columns, + tl.where(page_valid, score, -float("inf")), + mask=live & (columns < num_columns), + ) diff --git a/benchmarks/kernels/verify_sm70_batch_hc_runtime.py b/benchmarks/kernels/verify_sm70_batch_hc_runtime.py new file mode 100644 index 0000000000..fd8f8d793e --- /dev/null +++ b/benchmarks/kernels/verify_sm70_batch_hc_runtime.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real GatedResidual/TP channel route checks, not a model quality score.""" + +import argparse +import hashlib +import json +import os +from pathlib import Path +from types import SimpleNamespace as NS + +import torch +import torch.distributed as dist +from safetensors import safe_open + +from vllm import _custom_ops as ops +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + get_tp_group, + init_distributed_environment, + initialize_model_parallel, +) +from vllm.forward_context import set_forward_context +from vllm.models.qwen4_exp.common.hyperconnection import HyperConnectionConfig +from vllm.models.qwen4_exp.nvidia.hyperconnection import GatedResidual +from vllm.models.qwen4_exp.nvidia.sm70_batch_hc import prepare_sm70_batch_hc +from vllm.models.qwen4_exp.nvidia.sm70_fp16_gemv import Qwen38SM70FP16LinearMethod + + +@torch.inference_mode() +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model", type=Path, required=True) + p.add_argument("--out", type=Path, required=True) + args = p.parse_args() + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=4)) + with set_current_vllm_config(config): + init_distributed_environment(world_size=4, rank=rank, local_rank=rank) + initialize_model_parallel(tensor_model_parallel_size=4) + comm = get_tp_group().device_communicator + assert comm.sm70_hc_batch_comm is not None + assert comm.sm70_hc_batch_comm is not comm.ca_comm + assert ops.supports_sm70_qwen38_hc_batch() + assert ops.supports_sm70_qwen38_hc_up_mix_allgather() + channel = comm.sm70_hc_batch_comm + hc = GatedResidual( + HyperConnectionConfig( + hc_count=4, + hidden_size=2560, + hc_lowrank=320, + params_dtype=torch.float16, + hc_per_branch_norm=True, + ) + ).cuda() + hc.input_mix_weight_down_block_inject.quant_method = ( + Qwen38SM70FP16LinearMethod() + ) + hc._sm70_qwen38_fp16_fused_hc = True + index = json.loads((args.model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + prefix = "model.language_model.layers.0.attn_hyper_connection." + + def read_weight(suffix): + name = prefix + suffix + with safe_open(args.model / index[name], framework="pt", device="cpu") as f: + return f.get_tensor(name).half().cuda() + + down = read_weight("input_mix_weight_down.weight") + injection = read_weight("block_inject_weight.weight") + hc.input_mix_weight_down_block_inject.weight.copy_( + torch.cat((down, injection, down.new_zeros(12, 10240))) + ) + hc.input_mix_weight_up.weight.copy_(read_weight("input_mix_weight_up.weight")) + prepare_sm70_batch_hc(hc) + packed = hc._sm70_batch_hc_up + rows = [] + native_calls = [0] + original_down = ops.sm70_qwen38_hc_batch_down + + def tracked_down(*args): + native_calls[0] += 1 + return original_down(*args) + + ops.sm70_qwen38_hc_batch_down = tracked_down + try: + for m, prefill in ( + (1, False), + (4, False), + (8, False), + (16, False), + (4, True), + (17, False), + (32, True), + ): + metadata = {"attn": NS(max_query_len=m if prefill else 1)} + torch.manual_seed(7) + x = torch.randn(m, 10240, dtype=torch.float16, device="cuda") + dist.broadcast(x, 0) + with set_forward_context(metadata, config, num_tokens=m): + hc._sm70_batch_hc_up = None + reference = tuple(t.clone() for t in hc._project(x)) + hc._sm70_batch_hc_up = packed + before = native_calls[0] + actual = tuple(t.clone() for t in hc._project(x)) + hit = native_calls[0] > before + assert hit == (not prefill and 2 <= m <= 16) + if not hit: + for a, b in zip(actual, reference): + torch.testing.assert_close(a, b, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + # Warm every native/projection path before graph capture. + for _ in range(3): + hc._project(x) + torch.accelerator.synchronize() + with torch.cuda.graph(graph): + output = hc._project(x) + for cycle in range(8): + x.normal_().mul_(0.25 + cycle % 3) + dist.broadcast(x, 0) + expected = tuple(t.clone() for t in hc._project(x)) + for t in output: + t.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + for a, b in zip(output, expected): + torch.testing.assert_close(a, b, atol=0, rtol=0) + rows.append( + { + "rows": m, + "prefill": prefill, + "native_hit": hit, + "graph_eager_exact": True, + "reference_max_abs": [ + (a.float() - b.float()).abs().max().item() + for a, b in zip(actual, reference) + ], + } + ) + # An initial prefill trace must not bake away the opaque decode op. + compiled = torch.compile(hc._project, backend="eager", dynamic=True) + for m, prefill in ((32, True), (4, False), (1, False), (16, True)): + x = torch.randn(m, 10240, dtype=torch.float16, device="cuda") + dist.broadcast(x, 0) + with set_forward_context( + {"attn": NS(max_query_len=m if prefill else 1)}, + config, + num_tokens=m, + ): + expected = tuple(t.clone() for t in hc._project(x)) + before = native_calls[0] + result = compiled(x) + hit = native_calls[0] > before + assert hit == (not prefill and 2 <= m <= 16) + for a, b in zip(result, expected): + torch.testing.assert_close(a, b, atol=0, rtol=0) + pointer = packed.data_ptr() + prepare_sm70_batch_hc(hc) + assert hc._sm70_batch_hc_up.data_ptr() == pointer + gathered = [None] * 4 + dist.all_gather_object(gathered, rows, group=get_tp_group().cpu_group) + if rank == 0: + lib = Path(os.environ["VLLM_SM70_CUSTOM_AR_LIBRARY"]) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps( + { + "rows": gathered, + "native_sha256": hashlib.sha256( + lib.read_bytes() + ).hexdigest(), + "torch": torch.__version__, + "prefill_first_compile_passed": True, + "reload_preserved_pointer": True, + "M1_native_capabilities": True, + "scope": ( + "Real HC component, not full-model quality " + "or throughput" + ), + }, + indent=2, + ) + + "\n" + ) + finally: + ops.sm70_qwen38_hc_batch_down = original_down + hc._sm70_batch_hc_up = packed + torch.accelerator.synchronize() + destroy_model_parallel() + assert channel._ptr == 0 + destroy_distributed_environment() + + +if __name__ == "__main__": + main() diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index b8c193834a..c76c70ef51 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -8,6 +8,9 @@ #include #include "custom_all_reduce.cuh" +#if !defined(USE_ROCM) + #include "sm70_hc_batch.cuh" +#endif // Fake pointer type, must match fptr_t type in ops.h. // We use this type alias to indicate when pointers are passed in as int64_t. @@ -796,6 +799,24 @@ void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, } } +void sm70_qwen38_hc_batch_down(fptr_t ptr, torch::Tensor input, + torch::Tensor injection, torch::Tensor lora) { +#if defined(USE_ROCM) + TORCH_CHECK(false, "SM70 batch HC is unavailable on ROCm"); +#else + sm70_hc_batch::run(ptr, input, injection, lora); +#endif +} + +void sm70_qwen38_hc_batch_mix(fptr_t ptr, torch::Tensor gate, + torch::Tensor branches, torch::Tensor output) { +#if defined(USE_ROCM) + TORCH_CHECK(false, "SM70 batch HC is unavailable on ROCm"); +#else + sm70_hc_batch::run(ptr, gate, branches, output); +#endif +} + void sm70_qwen38_hc_down_allgather(fptr_t _fa, torch::Tensor& input, torch::Tensor& output) { #if defined(USE_ROCM) diff --git a/csrc/flashinfer_sm70/gdn_bridge.cuh b/csrc/flashinfer_sm70/gdn_bridge.cuh new file mode 100644 index 0000000000..f92d280de7 --- /dev/null +++ b/csrc/flashinfer_sm70/gdn_bridge.cuh @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Shared by wheel instantiations and isolated geometry microbenchmarks. +#include +#include +#include +#include +#include +#include + +namespace { +using namespace FI_GDN_IMPL_NAMESPACE; + +void run(torch::Tensor hidden, torch::Tensor weights, torch::Tensor qkv, + torch::Tensor conv_w, torch::Tensor conv_bias, torch::Tensor conv, + torch::Tensor A_log, torch::Tensor dt_bias, torch::Tensor state, + torch::Tensor indices, torch::Tensor output, torch::Tensor conv_out, + torch::Tensor partial) { + const c10::cuda::CUDAGuard guard(hidden.device()); + for (const auto& t : {hidden, weights, qkv, conv_w, conv_bias, conv, A_log, + dt_bias, state, indices, output, conv_out, partial}) + TORCH_CHECK(t.is_cuda() && t.device() == hidden.device()); + const auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0, "SM70 GDN adapter only"); + const int B = hidden.size(0); + TORCH_CHECK(B > 0 && B <= 64 && hidden.dim() == 2 && + hidden.size(1) == HIDDEN); + TORCH_CHECK(weights.sizes() == torch::IntArrayRef({HIDDEN, N_BA})); + TORCH_CHECK(qkv.sizes() == torch::IntArrayRef({B, QKV_DIM}) && + qkv.stride(1) == 1); + TORCH_CHECK(conv_w.sizes() == torch::IntArrayRef({QKV_DIM, CONV_WIDTH})); + TORCH_CHECK(conv_bias.numel() == 0 || conv_bias.numel() == QKV_DIM); + TORCH_CHECK(conv.dim() == 3 && conv.size(1) == QKV_DIM && conv.size(2) == 3); + TORCH_CHECK(state.dim() == 4 && state.size(0) == conv.size(0) && + state.size(1) == HV && state.size(2) == D && state.size(3) == D && + state.stride(1) == D * D && state.stride(2) == D && + state.stride(3) == 1); + TORCH_CHECK(A_log.numel() == HV && dt_bias.numel() == HV && + indices.numel() == B); + TORCH_CHECK(output.numel() == B * HV * D && conv_out.numel() == B * QKV_DIM); + TORCH_CHECK(partial.numel() == B * N_BA * GEMV_NSPLIT); + for (const auto& t : {hidden, weights, conv_w, conv_bias, A_log, dt_bias, + indices, output, conv_out, partial}) + TORCH_CHECK(t.is_contiguous()); + for (const auto& t : {hidden, weights, qkv, conv_w, conv_bias, conv, dt_bias, + output, conv_out}) + TORCH_CHECK(t.scalar_type() == at::kHalf); + for (const auto& t : {state, A_log, partial}) + TORCH_CHECK(t.scalar_type() == at::kFloat); + TORCH_CHECK(indices.scalar_type() == at::kInt); + TORCH_CHECK(reinterpret_cast(state.data_ptr()) % 16 == 0 && + state.stride(0) % 4 == 0); + const int block = 256; + auto stream = at::cuda::getCurrentCUDAStream(); +#if defined(FI_GDN_TWO_PHASE) && FI_GDN_TWO_PHASE + const int prep_grid = props->multiProcessorCount * 4; + const int state_grid = + (B * HV * D + ROWS_PER_WARP * 8 - 1) / (ROWS_PER_WARP * 8); + #define LAUNCH_PHASE(B1, PHASE, GRID) \ + gdn_fused_decode_kernel<<>>( \ + (const half*)hidden.data_ptr(), (const half*)weights.data_ptr(), \ + (const half*)qkv.data_ptr(), (const half*)conv_w.data_ptr(), \ + conv_bias.numel() ? (const half*)conv_bias.data_ptr() : nullptr, \ + (const half*)conv.data_ptr(), A_log.data_ptr(), \ + (const half*)dt_bias.data_ptr(), state.data_ptr(), \ + indices.data_ptr(), 1.f / sqrtf(float(D)), state.stride(0), \ + qkv.stride(0), conv.stride(0), conv.stride(1), conv.stride(2), \ + (half*)output.data_ptr(), (half*)conv.data_ptr(), \ + state.data_ptr(), partial.data_ptr(), \ + (half*)conv_out.data_ptr(), B) + if (B == 1) { + LAUNCH_PHASE(true, 1, prep_grid); + LAUNCH_PHASE(true, 2, state_grid); + } else { + LAUNCH_PHASE(false, 1, prep_grid); + LAUNCH_PHASE(false, 2, state_grid); + } + #undef LAUNCH_PHASE +#else + int occupancy = 0; + if (B == 1) { + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, gdn_fused_decode_kernel, block, 0)); + } else { + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, gdn_fused_decode_kernel, block, 0)); + } + TORCH_CHECK(occupancy > 0); + const int needed = (B * HV * D + ROWS_PER_WARP * 8 - 1) / (ROWS_PER_WARP * 8); + const int grid = std::min(needed, occupancy * props->multiProcessorCount); + TORCH_CHECK(props->cooperativeLaunch, + "GDN grid sync needs cooperative launch"); + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.stream = stream; + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeCooperative; + attribute.val.cooperative = 1; + config.attrs = &attribute; + config.numAttrs = 1; + #define LAUNCH(B1) \ + C10_CUDA_CHECK(cudaLaunchKernelEx( \ + &config, gdn_fused_decode_kernel, (const half*)hidden.data_ptr(), \ + (const half*)weights.data_ptr(), (const half*)qkv.data_ptr(), \ + (const half*)conv_w.data_ptr(), \ + conv_bias.numel() ? (const half*)conv_bias.data_ptr() : nullptr, \ + (const half*)conv.data_ptr(), A_log.data_ptr(), \ + (const half*)dt_bias.data_ptr(), state.data_ptr(), \ + indices.data_ptr(), 1.f / sqrtf(float(D)), state.stride(0), \ + qkv.stride(0), conv.stride(0), conv.stride(1), conv.stride(2), \ + (half*)output.data_ptr(), (half*)conv.data_ptr(), \ + state.data_ptr(), partial.data_ptr(), \ + (half*)conv_out.data_ptr(), B)) + if (B == 1) { + LAUNCH(true); + } else { + LAUNCH(false); + } + #undef LAUNCH +#endif + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace + +TORCH_LIBRARY_FRAGMENT(FI_GDN_TORCH_NAMESPACE, m) { + m.def( + "run(Tensor hidden, Tensor weights, Tensor qkv, Tensor conv_w, " + "Tensor conv_bias, Tensor(a!) conv, Tensor A_log, Tensor dt_bias, " + "Tensor(b!) state, Tensor indices, Tensor(c!) output, Tensor(d!) " + "conv_out, " + "Tensor(e!) partial) -> ()"); +} +TORCH_LIBRARY_IMPL(FI_GDN_TORCH_NAMESPACE, CUDA, m) { m.impl("run", &run); } diff --git a/csrc/flashinfer_sm70/gdn_h2560_q16_v48.cu b/csrc/flashinfer_sm70/gdn_h2560_q16_v48.cu new file mode 100644 index 0000000000..19d7e97bec --- /dev/null +++ b/csrc/flashinfer_sm70/gdn_h2560_q16_v48.cu @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Shape-specialized, not model-name or configured TP-degree dispatch. +#define FI_GDN_IMPL_NAMESPACE flashinfer::sm70::gdn::h2560_q16_v48 +#define FI_GDN_TORCH_NAMESPACE _C_flashinfer_gdn_sm70_h2560_q16_v48 +#define FI_GDN_HIDDEN 2560 +#define FI_GDN_N_BA 96 +#define FI_GDN_QKV_DIM 10240 +#define FI_GDN_H_Q 16 +#define FI_GDN_HV 48 +#define FI_GDN_D 128 +#define FI_GDN_CONV_WIDTH 4 +#define FI_GDN_CONV_STATE_LEN 3 +#include "gdn_bridge.cuh" diff --git a/csrc/flashinfer_sm70/gdn_h2560_q4_v12.cu b/csrc/flashinfer_sm70/gdn_h2560_q4_v12.cu new file mode 100644 index 0000000000..bb86337978 --- /dev/null +++ b/csrc/flashinfer_sm70/gdn_h2560_q4_v12.cu @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Shape-specialized, not model-name or configured TP-degree dispatch. +#define FI_GDN_IMPL_NAMESPACE flashinfer::sm70::gdn::h2560_q4_v12 +#define FI_GDN_TORCH_NAMESPACE _C_flashinfer_gdn_sm70_h2560_q4_v12 +#define FI_GDN_HIDDEN 2560 +#define FI_GDN_N_BA 24 +#define FI_GDN_QKV_DIM 2560 +#define FI_GDN_H_Q 4 +#define FI_GDN_HV 12 +#define FI_GDN_D 128 +#define FI_GDN_CONV_WIDTH 4 +#define FI_GDN_CONV_STATE_LEN 3 +#include "gdn_bridge.cuh" diff --git a/csrc/flashinfer_sm70/gdn_h2560_q8_v24.cu b/csrc/flashinfer_sm70/gdn_h2560_q8_v24.cu new file mode 100644 index 0000000000..5b50238225 --- /dev/null +++ b/csrc/flashinfer_sm70/gdn_h2560_q8_v24.cu @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Shape-specialized, not model-name or configured TP-degree dispatch. +#define FI_GDN_IMPL_NAMESPACE flashinfer::sm70::gdn::h2560_q8_v24 +#define FI_GDN_TORCH_NAMESPACE _C_flashinfer_gdn_sm70_h2560_q8_v24 +#define FI_GDN_HIDDEN 2560 +#define FI_GDN_N_BA 48 +#define FI_GDN_QKV_DIM 5120 +#define FI_GDN_H_Q 8 +#define FI_GDN_HV 24 +#define FI_GDN_D 128 +#define FI_GDN_CONV_WIDTH 4 +#define FI_GDN_CONV_STATE_LEN 3 +#include "gdn_bridge.cuh" diff --git a/csrc/flashinfer_sm70/qsa_mqa.cu b/csrc/flashinfer_sm70/qsa_mqa.cu new file mode 100644 index 0000000000..1e5b0a693d --- /dev/null +++ b/csrc/flashinfer_sm70/qsa_mqa.cu @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include +#include +#include +#include + +#include +#include + +namespace { +using namespace flashinfer::attention::sm70; + +void mqa(torch::Tensor q, torch::Tensor k, torch::Tensor table, + torch::Tensor requests, torch::Tensor positions, torch::Tensor lengths, + torch::Tensor logits, torch::Tensor visible, torch::Tensor schedule, + int64_t ratio, double divisor, int64_t workers, + std::optional task_visits) { + TORCH_CHECK(q.is_cuda(), "QSA MQA requires CUDA tensors"); + const c10::cuda::CUDAGuard guard(q.device()); + for (const auto& t : + {k, table, requests, positions, lengths, logits, visible, schedule}) + TORCH_CHECK(t.device() == q.device(), + "QSA MQA tensors must share a CUDA device"); + const auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0, + "QSA MQA native route requires SM70"); + TORCH_CHECK(q.dim() == 3 && q.size(0) > 0 && q.size(0) <= kMQAMaxRows && + q.size(1) > 0 && q.size(1) <= 16 && q.stride(2) == 1); + TORCH_CHECK(q.scalar_type() == at::kHalf && k.scalar_type() == at::kHalf); + TORCH_CHECK(k.dim() == 4 && k.size(0) > 0 && k.size(1) > 0 && + k.size(2) == 1 && k.size(3) == q.size(2) && k.stride(3) == 1); + TORCH_CHECK(q.size(2) == 64 || q.size(2) == 128 || q.size(2) == 256); + TORCH_CHECK(reinterpret_cast(q.data_ptr()) % 16 == 0 && + reinterpret_cast(k.data_ptr()) % 16 == 0 && + q.stride(0) % 8 == 0 && q.stride(1) % 8 == 0 && + k.stride(0) % 8 == 0 && k.stride(1) % 8 == 0, + "QSA MQA vector loads require aligned rows"); + TORCH_CHECK(table.dim() == 2 && table.size(0) > 0 && table.size(1) > 0 && + table.stride(1) == 1); + for (const auto& t : {table, requests, lengths, visible, schedule}) + TORCH_CHECK(t.scalar_type() == at::kInt, "QSA MQA metadata must be int32"); + TORCH_CHECK(positions.scalar_type() == at::kInt || + positions.scalar_type() == at::kLong, + "QSA positions must be int32 or int64"); + for (const auto& t : {requests, positions, visible}) + TORCH_CHECK(t.dim() == 1 && t.numel() == q.size(0) && t.is_contiguous()); + TORCH_CHECK(lengths.dim() == 1 && lengths.numel() == table.size(0) && + lengths.is_contiguous()); + TORCH_CHECK(logits.dim() == 2 && logits.size(0) == q.size(0) && + logits.size(1) > 0 && logits.size(1) <= INT32_MAX - kMQATile && + logits.scalar_type() == at::kFloat && logits.stride(1) == 1); + TORCH_CHECK(ratio > 0 && ratio <= INT32_MAX && std::isfinite(divisor) && + divisor > 0 && std::isfinite(static_cast(divisor)) && + static_cast(divisor) > 0); + TORCH_CHECK(workers > 0 && workers <= props->multiProcessorCount * 8); + TORCH_CHECK(schedule.dim() == 2 && schedule.size(0) == workers + 1 && + schedule.size(1) == 2 && schedule.is_contiguous()); + TORCH_CHECK(k.size(0) <= INT32_MAX && k.size(1) <= INT32_MAX && + table.size(0) <= INT32_MAX && table.size(1) <= INT32_MAX); + TORCH_CHECK(((logits.size(1) + kMQATile - 1) / kMQATile) * q.size(0) <= + INT32_MAX); + MQAParams p{}; + p.q = reinterpret_cast(q.data_ptr()); + p.k = reinterpret_cast(k.data_ptr()); + p.table = table.data_ptr(); + p.requests = requests.data_ptr(); + p.positions = positions.data_ptr(); + p.positions64 = positions.scalar_type() == at::kLong; + p.lengths = lengths.data_ptr(); + p.visible = visible.data_ptr(); + p.schedule = schedule.data_ptr(); + p.logits = logits.data_ptr(); + p.rows = q.size(0); + p.heads = q.size(1); + p.columns = logits.size(1); + p.pages = k.size(0); + p.page_size = k.size(1); + p.table_width = table.size(1); + p.num_requests = table.size(0); + p.ratio = ratio; + p.workers = workers; + p.divisor = divisor; + p.q_row = q.stride(0); + p.q_head = q.stride(1); + p.k_page = k.stride(0); + p.k_token = k.stride(1); + p.table_row = table.stride(0); + p.out_row = logits.stride(0); + if (task_visits.has_value()) { + const auto& visits = *task_visits; + TORCH_CHECK(visits.device() == q.device() && + visits.scalar_type() == at::kInt && visits.is_contiguous() && + visits.dim() == 2 && visits.size(0) == q.size(0) && + visits.size(1) == (logits.size(1) + kMQATile - 1) / kMQATile); + p.task_visits = visits.data_ptr(); + } + const auto stream = at::cuda::getCurrentCUDAStream(q.get_device()); + PlanMQA<<<1, 32, 0, stream>>>(p); +#define SCORE(D) \ + case D: \ + if (task_visits.has_value()) \ + ScoreMQA<<>>(p); \ + else \ + ScoreMQA<<>>(p); \ + break + switch (q.size(2)) { + SCORE(64); + SCORE(128); + SCORE(256); + } +#undef SCORE + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace + +TORCH_LIBRARY_FRAGMENT(_C_flashinfer_mqa_sm70, m) { + m.def( + "run(Tensor q, Tensor k, Tensor table, Tensor requests, Tensor " + "positions, " + "Tensor lengths, Tensor(a!) logits, Tensor(b!) visible, Tensor(c!) " + "schedule, " + "int ratio, float divisor, int workers, Tensor(d!)? task_visits=None) -> " + "()"); +} +TORCH_LIBRARY_IMPL(_C_flashinfer_mqa_sm70, CUDA, m) { m.impl("run", &mqa); } diff --git a/csrc/ops.h b/csrc/ops.h index 83d33c1def..d8d393d222 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -726,6 +726,10 @@ void sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( fptr_t reg_output_buffer, int64_t reg_buffer_sz_bytes, double epsilon); void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, torch::Tensor& out); +void sm70_qwen38_hc_batch_down(fptr_t ptr, torch::Tensor input, + torch::Tensor injection, torch::Tensor lora); +void sm70_qwen38_hc_batch_mix(fptr_t ptr, torch::Tensor gate, + torch::Tensor branches, torch::Tensor output); void sm70_qwen38_hc_down_allgather(fptr_t _fa, torch::Tensor& input, torch::Tensor& output); void sm70_qwen38_hc_gate_mix(fptr_t _fa, torch::Tensor& local_gate, diff --git a/csrc/sm70_hc_batch.cuh b/csrc/sm70_hc_batch.cuh new file mode 100644 index 0000000000..a08cfa673e --- /dev/null +++ b/csrc/sm70_hc_batch.cuh @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Uses a dedicated CustomAllreduce instance: never share its packet/epoch +// storage with auxiliary-stream collectives, or cross the owning DSO ABI. +#include "custom_all_reduce.cuh" +#include +#include +#include + +namespace sm70_hc_batch { +using namespace vllm; +using Pack = packed_t::P; + +// Match the existing Triton HC PTX, including div.full rather than replacing +// its division with a different approximate reciprocal or changing rounding. +__device__ __forceinline__ float full_div(float a, float b) { + float result; + asm("div.full.f32 %0, %1, %2;" : "=f"(result) : "f"(a), "f"(b)); + return result; +} + +__device__ __forceinline__ float sigmoid(float x) { + float exponent, denominator; + asm("mul.f32 %0, %1, 0fBFB8AA3B;" : "=f"(exponent) : "f"(x)); + asm("ex2.approx.f32 %0, %1;" : "=f"(exponent) : "f"(exponent)); + asm("add.f32 %0, %1, 0f3F800000;" : "=f"(denominator) : "f"(exponent)); + return full_div(1.0f, denominator); +} + +template +__global__ void gather_kernel(RankData buffers, const half* input, + const half* x, half* output, half* injection, + int rank, int rows) { + constexpr int local_cols = Mix ? 640 : 88; + constexpr int packs_per_row = local_cols / Pack::size; + constexpr int stride = kSm70Tp4PushAllreduceBytes / sizeof(Pack); + auto* local = + const_cast(reinterpret_cast(buffers.ptrs[rank])); + auto* epochs = reinterpret_cast(local); + const unsigned epoch = epochs[blockIdx.x]; + const int base = epoch * 4 * stride; + const int offset = blockIdx.x * blockDim.x + threadIdx.x; + if (offset < rows * packs_per_row) { + const int row = offset / packs_per_row; + const int col = offset % packs_per_row * Pack::size; + Pack value; + if constexpr (Mix) { + float accum[Pack::size] = {}; +#pragma unroll + for (int branch = 0; branch < 4; ++branch) { + const Pack g = *reinterpret_cast(input + row * 2560 + + branch * 640 + col); + const Pack v = *reinterpret_cast( + x + row * 10240 + branch * 2560 + rank * 640 + col); +#pragma unroll + for (int i = 0; i < Pack::size; ++i) + accum[i] = fmaf(sigmoid(__half2float(g.data[i])), + __half2float(v.data[i]), accum[i]); + } +#pragma unroll + for (int i = 0; i < Pack::size; ++i) { + value.data[i] = __float2half_rn(full_div(accum[i], 4.0f)); + // The unfused disjoint all-reduce adds positive zero to each rank's + // materialized FP16 result, canonicalizing signed zero. + if (__half2float(value.data[i]) == 0) + value.data[i] = __float2half_rn(0); + } + } else { + value = *reinterpret_cast(input + row * 88 + col); +#pragma unroll + for (int i = 0; i < Pack::size; ++i) { + float v = __half2float(value.data[i]); + if (v == 0) v = 0.0f; + if (col < 80) { + v = full_div(v, 4.0f); + v = __fmul_rn(v, sigmoid(v)); + } + value.data[i] = __float2half_rn(v); + } + } +#pragma unroll + for (int i = 0; i < Pack::size; ++i) + sm70_push_escape_sentinel(value.data[i]); +#pragma unroll + for (int peer = 0; peer < 4; ++peer) { + auto* destination = + const_cast(reinterpret_cast(buffers.ptrs[peer])); + destination += kSm70Tp4PushAllreduceSignalBytes + + (base + rank * stride) * sizeof(Pack); + sm70_push_store_volatile_16b(value, destination, offset); + } + Pack values[4]; + while (true) { + bool pending = false; +#pragma unroll + for (int peer = 0; peer < 4; ++peer) { + const void* source = local + kSm70Tp4PushAllreduceSignalBytes + + (base + peer * stride) * sizeof(Pack); + sm70_push_load_volatile_16b(values[peer], source, offset); +#pragma unroll + for (int i = 0; i < Pack::size; ++i) + pending |= sm70_push_is_sentinel(values[peer].data[i]); + } + if (!pending) break; + } +#pragma unroll + for (int peer = 0; peer < 4; ++peer) { + if constexpr (Mix) { + *reinterpret_cast(output + row * 2560 + peer * 640 + col) = + values[peer]; + } else { + if (col < 80) + *reinterpret_cast(output + row * 320 + peer * 80 + col) = + values[peer]; + else if (peer == 3) + *reinterpret_cast(injection + row * 4) = + *reinterpret_cast(&values[peer]); + } + } + Pack empty; +#pragma unroll + for (int i = 0; i < Pack::size; ++i) + *reinterpret_cast(&empty.data[i]) = + kSm70Tp4PushAllreduceSentinel; +#pragma unroll + for (int peer = 0; peer < 4; ++peer) { + void* source = local + kSm70Tp4PushAllreduceSignalBytes + + (base + peer * stride) * sizeof(Pack); + sm70_push_store_volatile_16b(empty, source, offset); + } + } + __syncthreads(); + if (threadIdx.x == 0) epochs[blockIdx.x] = (epoch + 1) % 2; +} + +template +void run(int64_t ptr, torch::Tensor input, torch::Tensor aux, + torch::Tensor output) { + const c10::cuda::CUDAGuard guard(input.device()); + auto* ca = reinterpret_cast(ptr); + TORCH_CHECK(ca && ca->world_size_ == 4 && ca->fully_connected_ && + ca->sm70_tp4_push_buffers_registered_ && + custom_allreduce_current_device_is_sm70()); + for (const auto& t : {input, aux, output}) + TORCH_CHECK(t.is_cuda() && t.is_contiguous() && + t.device() == input.device() && t.scalar_type() == at::kHalf && + t.dim() == 2); + const int rows = input.size(0); + TORCH_CHECK(rows > 0 && aux.size(0) == rows && output.size(0) == rows); + TORCH_CHECK(input.size(1) == (Mix ? 2560 : 88) && + aux.size(1) == (Mix ? 10240 : 4) && + output.size(1) == (Mix ? 2560 : 320)); + const int packs = rows * (Mix ? 80 : 11); + TORCH_CHECK(packs <= kSm70Tp4PushAllreduceBytes / sizeof(Pack)); + constexpr int threads = kSm70Tp4PushAllreduceThreads; + const auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); + gather_kernel<<<(packs + threads - 1) / threads, threads, 0, stream>>>( + ca->sm70_tp4_push_buffers_, + reinterpret_cast(input.data_ptr()), + Mix ? reinterpret_cast(aux.data_ptr()) : nullptr, + reinterpret_cast(output.data_ptr()), + Mix ? nullptr : reinterpret_cast(aux.data_ptr()), ca->rank_, rows); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +} // namespace sm70_hc_batch diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index f487996750..d96937ad51 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -1009,6 +1009,16 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { custom_ar.def( "all_reduce_sum2(int fa, Tensor inp_a, Tensor inp_b, Tensor! out) -> ()"); custom_ar.impl("all_reduce_sum2", torch::kCUDA, &all_reduce_sum2); + custom_ar.def( + "sm70_qwen38_hc_batch_down(int ptr, Tensor input, Tensor! injection, " + "Tensor! lora) -> ()"); + custom_ar.impl("sm70_qwen38_hc_batch_down", torch::kCUDA, + &sm70_qwen38_hc_batch_down); + custom_ar.def( + "sm70_qwen38_hc_batch_mix(int ptr, Tensor gate, Tensor branches, " + "Tensor! output) -> ()"); + custom_ar.impl("sm70_qwen38_hc_batch_mix", torch::kCUDA, + &sm70_qwen38_hc_batch_mix); custom_ar.def( "sm70_qwen38_hc_down_allgather(int fa, Tensor inp, Tensor! out) -> ()"); custom_ar.impl("sm70_qwen38_hc_down_allgather", torch::kCUDA, diff --git a/docs/design/sm70_flashinfer_batch_integration.md b/docs/design/sm70_flashinfer_batch_integration.md new file mode 100644 index 0000000000..398ffe3e01 --- /dev/null +++ b/docs/design/sm70_flashinfer_batch_integration.md @@ -0,0 +1,926 @@ +# FlashInfer SM70 batch integration probe + +Purpose: combine operator-screen winners, then perform one consolidated +control/candidate model comparison. No release default is changed. + +Integration base: `onecat/main` at `755baae1d075ee04fa9096b23fc0225b23589a86`. +This owned branch combines #504 at `5a049230cd` (existing batch HC/projection +and TP baseline), #513 at `17fa36e35a` (native FlashInfer QSA), and #515 at +`9ea54c19f4` (fused gate/conv/GDN prototype). Combination commit `f8e822c335`. +The runtime bridge is new, not a duplicate of those component-only PRs. + +## Decisions and contracts + +- Retain existing M1, HC projection/TP and HC norm paths. Both FlashInfer HC + staging variants passed numerical tests but lost to the existing Triton + component. On GPU 4, FP16 B16 baseline 3.164 us, register variant 3.779 us; + no HC replacement is admitted. +- `VLLM_SM70_FLASHINFER_BATCH=0` by default. Explicit prebuilt GDN and QSA + library paths are used only in the probe; no compiler or JIT during capture. +- GDN preserves FP16 activations/projections and FP32 recurrent state. Fuse at + the existing opaque input/core boundary to remove, not duplicate, the BA + GEMM. Ordinary QKVZ projection and output norm/projection remain unchanged. +- GDN dispatch uses actual non-spec, uniform-decode metadata and compatible + tensor geometry (2..64 rows). M1, larger shapes, prefill/mixed/MTP and other + state/activation layouts fall back. No global TP, scheduler budget or KV + dtype binding. Initial model validation is no-MTP; spec engines are not armed. +- QSA accepts the proven small-batch FP16 sparse-page family (4..16 rows), + preserving index order, repeats, invalid-page masking and the FP16 output + boundary before the existing output gate. E4M3 and other shapes fall back. +- Scratch allocations are call-local, never a cross-stream global mutable + workspace. The zero page is persistent/read-only. Derived-weight reloads + copy in place and reject changed geometry rather than replacing graph pointers. +- Do not change selector, expert routing, sampling, state precision or output + length to manufacture a throughput gain. + +## Validation and measurement + +The component worklogs retain micro timings and reference/source/binary hashes. +On GPU 4, #515's updated suite passed **20 tests**, including register-HC +cases (66.92 s including rebuild). GDN targeted memcheck passed one B16 case +with **0 errors**; racecheck passed one B4 case with **0 errors/0 warnings**. +These are targeted operator checks, not whole-engine sanitizer coverage. + +Initial bridge suite: **21 passed** on GPU 4 (11.67 s), including actual QSA +output gating, GDN input/core state evolution and eager/graph equality. A CPU +no-op initially imported the GDN module before its guard; fixed by guarding +first, so unsupported devices do not load native libraries. + +The next consolidated run uses the unchanged 8192-input/256-output fixed-live- +width benchmark, C1/4/8/16, no MTP, NVFP4 FlashNext, TP4, 256K capacity, prefix +cache + Mamba align, FP16 QSA KV, FP32 GDN state. CUDA Graph, GPU 4--7; both +arms use identical baseline flags and NUMA-node-1 CPU affinity. Keep automatic +GPU boost; do not compare absolute timings with prior GPU 0--3 runs. + +Use full EngineCore timestamp intervals, not client receive blocking time. +The fixed denominator remains 70 tok/s: C4/C8/C16 targets 238/420/728 tok/s. +The deterministic performance workload remains unchanged and is **not** quality +evidence. Its historical forced-length policy does not apply to quality cases. + +Quality: reuse the fixed GSM8K first-16 health screen (temperature 1, top-k20, +top-p.95, natural EOS, max16384), plus the retained 80-case BFCL/schema manifest +as an offline generation/parser screen. Explicitly distinguish offline model +outputs from HTTP/SSE transport testing and official benchmark scores. Report +scores, incomplete outputs and regressions; do not claim broad quality from +a short smoke. No candidate model result has been produced yet. + +Artifacts are task-private under this worktree's `.artifacts/`: native build +caches, `bridge-tests-gpu4-v2.log`, `env.sh`, `run-e2e.sh`, `run-model.py` and +`e2e/`. Frozen native vLLM/Flash-V100/HC libraries are pinned reference binaries; +this is not a freshly built release wheel. No shared checkout is modified. + +The user freed GPU 4--7. The old **local** `1cat-qwen38-flash-next.service` +was stopped and disabled; remote delivery services were not changed. Stop all +task-owned engine workers after each comparison. Keep this PR Draft until +the model and quality gates pass. AI-assisted (Codex), DCO sign-off and human +review required before merge. + +## First model comparison and compiled-boundary correction (September 6) + +Draft integration PR: #523. Both arms used the frozen speed runner and +environment above on GPU 4--7. The control reproduced the historical baseline: + +| Concurrency | Control tok/s | First candidate tok/s | +|---|---:|---:| +| 1 | 87.750 | 87.864 | +| 4 | 217.026 | 219.474 | +| 8 | 364.521 | 375.581 | +| 16 | 587.789 | 612.823 | + +These first candidate numbers are **QSA only**, not combined GDN + QSA. +Although all four ranks prepared 36 GDN layers, no fused GDN route was hit. +The `1 < num_tokens <= 64` test was outside the opaque input/core boundary: +the initial large-prefill trace specialized that branch away. Move all shape +and metadata selection inside the runtime op. Preserve the existing fused +FP16 input projection on fallback, including M1, instead of substituting +separate QKVZ/BA GEMMs. The boundary has an extra Z output copy; its M1 cost +and full-model behavior still need the corrected combined comparison. + +Added an export regression using the actual boundary predicate: one large +example must support the full 1..2048 row range. Reinstating the legacy shape +guard demonstrably fails export's shape constraints. The corrected predicate +and original-projection fallback checks pass: **24 CPU tests**, 5.63 seconds. +The first test attempt had a missing CPU-only op stub; fixed the fixture, +not the model implementation. The native operator binaries are unchanged. + +Quality is **not admitted**: GSM8K is 15/16 with no truncation in both arms, +but the fixed offline BFCL screen is **52/64 -> 50/64** (irrelevance 13 -> 11). +JSON Schema remains 16/16. Retain all per-case records; a small stochastic +screen cannot establish noninferiority, and this negative signal cannot be +ignored in exchange for C16's 4.3% gain. Do not enable by default. + +Artifacts: `.artifacts/e2e/control-*`, +`.artifacts/e2e/qsa-only-190e5f0225/candidate-*`, +`.artifacts/control-e2e-v1.log`, `.artifacts/candidate-e2e-v2.log`, +and `.artifacts/compiled-boundary-cpu-v2.log`. +Candidate attempt v1 exited at the GPU lease gate (75), before a model run; +it produced no performance result. Per-device external campaign leases are +also respected by the launcher now. All first-comparison workers exited; +GPU 4 was subsequently acquired by an unrelated sanitizer job, which is not +terminated by this task. + +The next combined attempt at `45711b5e21` passed compilation but failed during +FULL graph capture, before any quality or speed request. AOT compilation +preceded KV allocation, so the input/core op received empty 1-D state +placeholders. The new bridge tried to transpose them. The existing standard +recurrent core already resolves this case from the scheduler-bound layer +cache; reuse that contract inside the new runtime bridge. Explicit nonempty +state inputs must never be replaced. Add both placeholder and explicit-cache +graph/state tests, including a different layer cache to catch accidental +replacement, and guard malformed/unbound states before transpose. + +Failure log: `.artifacts/candidate-e2e-v3.log`. It is a startup failure, not a +performance or quality result. Its workers exited. Run the expanded GPU +bridge suite under the same GPU lease before the next model launch. + +While waiting for the unrelated MTP job, candidate v4 exited at the lease +gate. Candidate v5's preflight rejected the test launcher's four-visible-GPU +setting (the component suite intentionally requires one visible GPU); fixed +the launcher to expose GPU 4 only to pytest, then 4--7 to the model. Candidate +v6 passed all four GPU bridge cases, including both state-binding cases, but +the default-off test caught a test-order contamination: monkeypatching an +`envs` module attribute restored a materialized `True` attribute, shadowing +dynamic environment lookup. Change the test fixture to patch the environment +variable with the cache disabled instead. Neither v5 nor v6 launched a model. + +## Corrected combined result: faster, not admitted + +Model source `f45c673898` (runtime fix `2273c24d7d`), log +`.artifacts/candidate-e2e-v7.log`. Expanded preflight: **29 passed**, including +four GPU cases, 8.30 seconds. Then one corrected combined model run on the +same GPU 4--7 and unchanged workload. Worker logs confirm native fused GDN at +B2/4/8/16 and native FlashInfer QSA at B4/8/16. All ranks prepared 36 GDN +layers. FULL graph capture completed; no MTP or state-precision change. + +| Concurrency | Control tok/s | GDN + QSA tok/s | Gain | Fixed-70 efficiency | Target | +|---|---:|---:|---:|---:|---:| +| 1 | 87.750 | 88.556 | +0.92% | n/a | n/a | +| 4 | 217.026 | 232.169 | +6.98% | 82.92% | 85% | +| 8 | 364.521 | 390.046 | +7.00% | 69.65% | 75% | +| 16 | 587.789 | 630.996 | +7.35% | 56.34% | 65% | + +Complete engine-step means (ms): C4 **18.431 -> 17.229**, C8 +**21.947 -> 20.510**, C16 **27.221 -> 25.357**. These are unprofiled complete +decode intervals, not kernel sums, API receive blocking time or MTP rounds. +All three throughput targets remain unmet. This is one campaign, not a +cross-run confidence/stability study. Do not generalize its speed to other +contexts, model quantization, concurrency or serving modes. + +| Fixed quality subset | Control | QSA only | GDN + QSA | +|---|---:|---:|---:| +| GSM8K health screen | 15/16 | 15/16 | 15/16 | +| BFCL simple Python | 14/16 | 14/16 | 13/16 | +| BFCL parallel | 11/16 | 11/16 | 11/16 | +| BFCL multiple | 14/16 | 14/16 | 14/16 | +| BFCL irrelevance | 13/16 | 11/16 | 11/16 | +| BFCL total | 52/64 | 50/64 | 49/64 | +| JSON Schema | 16/16 | 16/16 | 16/16 | + +GSM8K and tool/schema outputs have no length truncations. Quality remains +**unadmitted**: five BFCL control successes become failures and two failures +become successes in the combined run. Identical prompts are retained. The +two irrelevance regressions recur in both candidates (`irrelevance_7` and +`irrelevance_10`); this is a localization lead, not proof that a specific +kernel or rounding operation is causal. Small stochastic screens cannot +establish either broad degradation or noninferiority. No HTTP/SSE transport, +coding, long-context quality or PPL admission is claimed here. + +Decision: keep the integration default **off**, PR #523 **Draft**, existing +HC norm and production source defaults unchanged. Do not count these faster +numbers as accepted production performance. Next isolate the QSA arithmetic +and output-gate drift on retained real trajectories, then GDN independently; +do not rerun until a narrower check can distinguish a cause. Preserve this +manifest and sampling, and do not select favorable seeds to clear the gate. + +Final artifacts: `.artifacts/e2e/candidate-{speed,gsm8k,offline-tools}.json` +and the corresponding control files, environment snapshots and GPU metadata. +The E2E FlashQLA reference binary SHA256 is +`c8bd7650444ec56cfe2576c044d8f5f438b0a352877064bbb51ac0510dc2ea2c`; +it is the same pinned library as the frozen HC model baseline. Component +native libraries remain task-private prebuilt probes, not a release wheel. + +All task model workers exited normally. Cleanup check: GPU 4--7 each **7 MiB**, +no compute processes; the old local API unit is inactive/disabled (MainPID 0). +Unrelated GPU 0--3 processes and remote services were not touched. Engine +shutdown logs also retain Python resource-tracker shared-memory cleanup +warnings seen in the control; do not mislabel these as persistent GPU usage. + +## Quality root-cause follow-up: same-input shadows + +Continue the owned #523 scope at `25fd594d9c`; do not duplicate component PRs +or change production precision to fit the small task-score sample. Inspecting +raw responses excludes a parser-only explanation: new failures contain no +call or an actual irrelevant call. However, the QSA-only comparison changes +the **first generated token in 11/80 cases**, including both recurring +irrelevance failures, despite identical prompts. Prefill/admission variation +must be separated from decode arithmetic before assigning causality. + +An artifact-only eager shadow run follows the original path for all actual +outputs/state updates. Rank 0 also computes new QSA on identical Q/K/V and +new GDN on cloned indexed initial states. Eight retained prompts and up to +48 output tokens are diagnostic input acquisition, **not quality scores**. +Other ranks do not run shadows/collectives. No production sampler, weights, +state dtype or output-selection policy is changed. + +The first attempt stopped at callable-RPC serialization restrictions before +requests. Replaced that diagnostic trigger with recognition of the retained +request seeds; no insecure-serialization flag is enabled. Successful run: +`.artifacts/quality-shadow-v2.log`, with records/captured QSA inputs under +`.artifacts/quality-shadow-v2/`. All workers exited and GPU 4--7 returned to +7 MiB. An unrelated job owned the cards while this task waited for the lease. + +Collected 144 GDN and 48 QSA comparisons (B8/B6/B5), plus 50 sampler records: + +- GDN Z projection and convolution state: **bitwise equal** in all records. +- Maximum GDN output relative L2: `4.8673613e-5`; maximum recurrent-state + relative L2: `1.7772339e-5`. Early B8 records alone had state errors around + `1e-8`; do not report that as the worst of the completed run. +- Maximum QSA candidate-vs-FP64-oracle relative L2: `1.9871883e-5` versus + reference-vs-oracle `1.8209650e-4`. The candidate is closer to the oracle + in **all 48** matched comparisons. Preserve ordered/duplicate selections + and the FP16-output/FP32-gate boundary in the oracle. +- No future or sequence-out-of-bounds selected positions, sampled-logit NaNs, + or positive infinities in this diagnostic scope. + +This does not establish whole-model noninferiority, CUDA Graph long-history +quality, or a causal explanation for the 52 -> 49 task-score result. It does +argue against blindly reducing new QSA precision or blaming GDN state +corruption without further evidence. Next isolate schedule/prefill and +sampling-trajectory confounders using matched conditional probabilities and +an unchanged-data A/A control. Preserve the old adverse scores. + +### Fixed-cohort and conditional-probability diagnostic + +Artifacts: `.artifacts/run-quality-cohort.py`, +`.artifacts/run-quality-cohort.sh`, `.artifacts/compare-quality-cohort.py`, +and `.artifacts/quality-cohort-v1/`. This retains all original 80 rendered +prompts, per-case seeds, 16K maximum output, natural EOS, schemas, parsers and +scorers. It changes only admission for a diagnostic: five fixed groups of +16, cold prefix reset per group, pause scheduling before enqueue and resume +after all group members are queued. The original continuous-admission +negative screen is NOT replaced by this diagnostic. + +Two aborted diagnostic attempts are retained. The first teacher manifest +was corrupted by truncated tool output; validate JSON before constructing +the engine (`teacher-manifest-valid.json` is the valid input). The second +completed the first group but failed in artifact-side output association: +`enqueue()` returns randomized internal IDs whereas `RequestOutput` uses +external IDs. The corrected runner snapshots the engine's own mapping +while scheduling is paused; a CPU fake-engine test covers this association. +Do not disable request-ID randomization or insecure-RPC serialization guards. + +Successful control log: `quality-cohort-v1/control-ids-fixed.log`. No new +FlashInfer operators are selected in this control. The two natural-EOS +repeats score BFCL **51/64 and 52/64**, JSON Schema **16/16 both**, with no +length truncations. **22/80 full outputs and 8/80 first tokens differ even +within this unchanged control engine.** This is direct evidence that a +single matched-seed score is not a deterministic attribution test. It does +not establish whether residual variation comes from prefill cohorts after +admission, numerical reduction/order, state lifetime or another mechanism. + +The separate teacher probe follows 16 retained historical-control sequences +(2,225 tokens). It is diagnostic conditional NLL, not a task score, PPL +benchmark or speed measurement. All four control ranks report identical +349-step seed/position cohort traces, and CPU/GPU request seeds agree. +25 CPU routing/compiled-boundary tests pass. + +Candidate log: `quality-cohort-v1/candidate.log`; comparison: +`quality-cohort-v1/comparison.txt`. Both native FlashInfer routes are observed +in CUDA Graph capture. Natural results (not release admission): + +| Path | BFCL repeat 1 | BFCL repeat 2 | JSON Schema both | Truncated | +| --- | ---: | ---: | ---: | ---: | +| Unchanged control | 51/64 | 52/64 | 16/16 | 0 | +| GDN + QSA candidate | 50/64 | 51/64 | 16/16 | 0 | + +The candidate A/A changes 23/80 full outputs and 8/80 first tokens. Preserve +these adverse comparisons and the original 52 -> 49 result; repeated tests +are neither independent extra dataset items nor grounds for picking the +best score. The original pair has five newly failed and two improved BFCL +cases; its exact paired two-sided test gives p=0.453125. Lack of significance +is NOT evidence of noninferiority. + +The candidate-minus-control conditional NLL mean is `+0.00310631` over 2,225 +tokens (positive is worse on the retained reference continuation); maximum +absolute delta is `1.93978739`. For first tokens alone, mean absolute delta +is `0.54163724`, versus `0.00885871` at offsets >=64. The four-rank +seed/position traces match exactly across arms. This points toward the +prefill/first-token portion, but needs a conditional A/A floor before +attributing the change to the new compiled boundary or kernels. + +### Conditional control A/A floor and current decision + +An additional control-only load runs the exact same 16 continuations twice, +with no intervening natural generation and no new FlashInfer operators. +Runner: `.artifacts/run-quality-teacher-aa.py`; command: +`FI_TEACHER_AA=1 bash .artifacts/run-quality-cohort.sh control`. +Log: `.artifacts/quality-teacher-aa-v1.log`; raw probabilities and four-rank +metadata: `.artifacts/quality-teacher-aa-v1/`. All 2,225 forced tokens match; +all four 698-step traces match, and their 349-step halves match each other. + +| Diagnostic | Mean NLL delta, all tokens | First-token mean absolute delta | Maximum absolute delta | +| --- | ---: | ---: | ---: | +| Candidate vs control | +0.00310631 | 0.54163724 | 1.93978739 | +| Control repeat 2 vs repeat 1 | +0.00435804 | 0.80130252 | 3.39355850 | + +Thus the unchanged control itself exhibits conditional-probability variation +at least as large in these aggregate measures as the candidate comparison. +This prevents assigning the task-score loss specifically to FlashInfer. +It does NOT prove all variation has the same cause, that state handling is +correct, or that the candidate is quality-equivalent. Investigate shared +prefill/first-token computation and state initialization/reuse first. Next +use a fixed-trajectory prefill boundary/state capture to locate the earliest +divergence; do not randomly change sampling, precision or decode fusion. +The AWQ FP16 atomic weighted epilogue found by source search is not on this +NVFP4 prefill path; do not claim it as the cause without a route-hit. + +No production math, kernels, sampler or defaults changed in this follow-up. +The earlier ~7% unprofiled speed gain remains an experimental result for the +same implementation, not a new speed measurement or quality admission. +Keep #523 Draft and the parent switch default-off. All diagnostic workers +exited; GPU 4--7 returned to 7 MiB each, with no compute processes. No remote +API was changed. Loader/shared-memory teardown warnings are retained in the +logs; no persistent GPU allocation remained. + +### Device-planned MQA implementation (2026-09-06, not quality-admitted) + +Upstream head was rechecked once and remains +`6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f`. The SM70 adaptation follows the +official attention-score scheduler's live-tile prefix scan and contiguous +balanced worker assignment; it does not compile the SM100 kernel for Volta. +The new scorer uses FP16 WMMA with FP32 accumulation, 64-column tiles, +vectorized eight-half loads and padded shared leading dimensions. A one-warp +device planner runs on every invocation, including graph replay. Caller-owned +schedule/output buffers retain stable graph addresses. No host length readback +or capacity-sized empty CTA grid is required. + +The native implementation supports int32 and int64 positions. Initial runtime +admission is the measured H4/D128, R4--16 geometry, independent of model name, +TP degree and configured maximum batch. Other shapes fall back locally. +The parent `VLLM_SM70_FLASHINFER_BATCH` remains **default off**. Runtime imports +the packaged `_sm70_flashinfer_C` fragment, not a development JIT library. +CMake, setup.py, manifest and source/license notice now include that fragment. +This is not yet a full built-and-installed wheel acceptance test. + +Formal CMake component artifact: +`.artifacts/wheel-native-stage/vllm/_sm70_flashinfer_C.abi3.so`, SHA256 +`1b2a2336d6dc22d207111008ed8c83cc97eac67f6ff0917dcb7e1ad456634edf`. +It passes **31 GPU tests** (`mqa-wheel-tests-v1.log`): FP64 oracle, both position +dtypes including overflow boundaries, strided layouts, graph length changes, +empty/padded rows, invalid pages, two-stream separate workspaces, and an audit +that counts every logical tile exactly once. CPU routing/compiled-boundary +tests pass **29/29** (`fi-cpu-v2.log`). + +The five alternating A/B blocks in `mqa-selector-v1.json` include the existing +exact top-k and page/index expansion, but **exclude final sparse attention and +the model**. Inputs are synthetic at real indexer geometry; this is not a +captured-activation or end-to-end result. Eight changing graph replays per case +produce identical selected indices. + +| Rows | Context | Existing index chain, us | New index chain, us | Reduction | +| --- | --- | ---: | ---: | ---: | +| 4 | 8K | 49.17248 | 40.56064 | 17.51% | +| 8 | 8K | 61.15328 | 42.27072 | 30.88% | +| 16 | 8K | 92.03712 | 53.99552 | 41.33% | +| 4 | 64K | 161.86369 | 124.08832 | 23.34% | +| 8 | 64K | 216.81152 | 149.88288 | 30.87% | +| 16 | 64K | 367.85152 | 237.88544 | 35.33% | + +Within-run paired log-ratio 95% intervals have positive lower bounds for these +six cases; they are not cross-run or model-quality confidence intervals. +`sm70_paired_stats.py` requires five paired blocks and uses Student-t(df=4). +Its seven CPU tests pass. Do not multiply layer service savings and report +the sum as measured end-to-end gain. + +Preserved rejected attempts: v1 had an incorrect lexicographic worker-end +condition, redundantly executing later rows while writing the same values; +numeric equality alone missed it. The tile-visit audit now detects this. +Corrected scalar/shared-unskewed v2 remained slower; vector loads and shared +padding were necessary. Logs `mqa-screen-v1.*`, `mqa-screen-v2.*` and +`mqa-native-trace-v2.nsys-rep` retain the failures. NCU counters were denied +(`ERR_NVGPUCTRPERM`); no counter-based occupancy or bandwidth claim is made. + +GDN's separate BA/conv prepare and state-update prototype preserves all outputs, +BA partials, conv and recurrent states exactly across 256 changing-history +steps at R1/4/8/16/32/64 (`gdn-phases-v1.*`). Its first timing run accidentally +retained a padded last row; those timings are invalid for full-width claims. +The benchmark now restores all live rows before timing; rerun is pending. +No two-phase GDN runtime route has been admitted. The HC 96-module full-chain +benchmark has been extended to B4/8/16 and five alternating blocks; GPU +validation is pending. MoE device compaction and TP overlap remain unimplemented. + +### First-prefill boundary localization (control only) + +Artifacts: `.artifacts/run-prefill-boundary.py`, corresponding shell/bootstrap, +`prefill-boundary-v2.log`, and `prefill-boundary-v1/boundary-rank*-step*.pt`. +One control-only engine, no FlashInfer experiment, executes two cold passes +over the same 16 teacher prompts. Maximum output is one forced reference token +for diagnosis only, not the registered quality battery or a performance run. + +On **all four ranks**, step 0 versus step 4 has bitwise-identical input IDs, +positions, query boundaries, request seeds and PLE ngram context. All six +requests are cold prefills. Layer 0 and layer 1 GDN projection inputs, conv +outputs, zero initial states and recurrent outputs are all bitwise identical. +Physical state IDs differ as expected after fresh allocation. Nevertheless, +the final hidden tensor [2048,2560] differs on 4,826,852 elements, max absolute +delta 31.966796875 and relative L2 0.2170600146; first-token logprob variation +also reproduces. This is **not** evidence of quality safety or a proven cause. +It localizes the first divergence downstream of the captured early GDN work; +inspect the next layers, first QSA/PLE, HC and MoE boundaries next. Do not blame +the new native scorer, which was disabled throughout this run. + +All diagnostic model workers shut down normally; no remote service changed. +The adverse BFCL results remain unresolved and the new whole-model performance, +full wheel, expanded quality battery and default promotion are still pending. + +### Root-cause closure and reuse of existing PR #494 + +The follow-up capture `prefill-boundary-later-v1/` (log +`prefill-boundary-later-v2.log`) localizes the first difference to layer 3 +QSA on all four ranks. Its input hidden states, Q/K/V, gate, positions, +selected token IDs, and effective logical K/V read back after the cache update +are bitwise identical. The first QSA output alone differs (rank 0 relative +L2 `9.7149867e-5`, max absolute `0.00048828125`), followed by progressively +larger downstream differences; final hidden relative L2 is `0.22358379`. +CPU FP64 causal attention on the captured cold-prefix keys confirms both +outputs have small local numerical error, not corrupted K/V. This local +oracle does not establish model-quality noninferiority. + +An isolated diagnostic alternative used logical request/page hash identities +and ordered collision resolution. Twelve relocations of the *same real Q/K/V* +changed 162,113--303,033 output elements in the old planner, versus zero in the +diagnostic alternative. The latter preserves grouped attention and FP16/FP32 +types but costs about 451 vs 434 us for captured 2048-row +planner+attention+gate; this is a stability repair, not a speed win. +Six planner tests plus 31 MQA tests pass (37 total); its six planner cases pass +memcheck with zero errors. The first sanitizer attempt lacked an injection +library, so it is not counted. Successful log: `canonical-memcheck-v2.log`. +Proof binary: `.artifacts/mqa-plus-canonical-proof/vllm/_sm70_flashinfer_C.abi3.so`, +SHA256 `23a202e5f4d7c5a9bccb73e4eba577b85fe090ead5e672ac039a1223adc521c1`. + +Changing only that planner in one additional cold-prefill engine eliminates +the observed A/A instability: all four ranks, all four scheduler-step pairs, +and their final hidden tensors match bitwise. All 16 first-token logprobs match +exactly (previous diagnostic mean absolute delta 0.60679578, maximum 1.85863316). +Artifacts: `prefill-boundary-canonical-v1/`, its launch log and +`prefill-boundary-canonical-compare-r*.log`. This is a causal diagnostic for +the allocation-order defect, **not a new task score or a validation of #494's +binary**. + +The subsequent overlap review found existing open **#494**, reviewed source +`5fa8a605dab12cc9ee15459d9ac6b88d95c7be3a`, already fixes this same defect. +It additionally preserves cross-request physical-page deduplication and fixes +the separate XQA tail. Reuse that reviewed patch rather than publish a competing +implementation. The new alternative planner, binding and tests were removed +from build/runtime/source delivery and retained only in +`.artifacts/canonical-prototype-source/`. Its results above are independent +NVFP4 root-cause evidence, not a claim of authorship of the existing repair. +The frozen performance binary did not include #494. Integration/build/testing +of #494 is the next dependency step; do not mix its forthcoming results with +the retired alternative's proof. + +The corrected full-width GDN phase screen is complete on GPU 0 +(`gdn-phases-v2.json`). All 256 history steps and output/conv/state/BA partials +remain exact at B1/4/8/16/32/64. Splitting the phases regresses B4/8/16 by +9.43%/15.80%/3.81%; keep it out of their runtime. B1/B32/B64 improvements are +4.22%/0.51%/3.91%, not admissions or justification to replace the existing M1 +route. No new end-to-end score has been measured. CPU routing, statistics and +compiled-boundary tests now pass 36/36 (`fi-cpu-v4.log`). + +### Adopted #494 and reduced its stable-plan overhead + +The reviewed dependency is cherry-picked as `3e0f7a40c1`, retaining the +original author and sign-offs. Its rebuilt Flash-V100 library passes 78 +QSA/MQA tests in this integration (`pr494-gates-v1.log`), including the 31 +MQA tests above. These are overlapping suites, not 78 additional MQA tests. + +Real layer-3 NVFP4 Q/K/V replay also confirms #494 allocation invariance. +Unlike the retired diagnostic prototype, its fixed 8192-entry radix sort +costs 653.26 us for captured 2048-row planner + attention + output gate, +versus 433.00 us in the unsafe frozen reference. Neither timing includes +the QSA indexer or the complete model. Artifact: `pr494-real-replay-v1.json`. +Reference #494 library SHA256: +`a99cce1f5fe32d61ef42525435401c4d24ddff53894eafe64cc534efa937ea23`. + +The follow-up retains #494's physical-page union, minimum logical owner, +category padding and exact sorted order. After loading the hash entries into +registers, a device block scan compacts valid entries, then chooses a +512/1024/2048/4096/8192-entry sort. All shared memory is reused within the +original 96-KiB budget; no host length readback or captured pointer changes +are introduced. The planner remains 128 registers/thread with zero local +spills according to `cuobjdump` (static resource data, not measured occupancy). + +Current adaptive binary SHA256: +`1c6fc18851e551950885c48ba0d108be919efcb6f4e0f6ca27e69ecb6e8fcd33`. +Validation: + +- 80 tests pass, including new empty-to-maximum live-union transitions across + every sort-size boundary inside one captured graph and in reverse order. +- Twelve physical relocations of the captured real input are bitwise exact + against #494's fixed-sort output, including gate materialization. +- Five alternating A/B blocks: **653.21 -> 458.50 us**, 29.81% reduction, + within-run paired 95% interval **[29.76%, 29.85%]**. This recovers most of + the stability-fix overhead; it is not an endpoint speed claim and remains + slower than the allocator-dependent unsafe reference. +- Artifacts: `pr494-adaptive-gates-v1.log`, + `pr494-adaptive-real-replay-v1.{json,log}` and both build logs/libraries. + The two changing-sort-size graph cases also pass Compute Sanitizer memcheck + with zero errors (`pr494-adaptive-memcheck-v1.log`); wider prefill performance + and racecheck remain pending. + +### GDN shared-parameter screen: reject for C8/C16 + +An isolated compile-time variant computes the unchanged gate reduction and +Q/K normalization once per CTA, sharing the results among its eight warps. +It preserves all FP16 round trips and FP32 state and introduces no global +workspace. The default kernel does not select this variant. + +All B1/2/4/8/16/32/64 cases preserve output, conv, recurrent state and BA +partials exactly over 256 changing-input graph steps, including slot +permutation and negative padding. Full-width timing is restored after the +padding history. Actual checkpoint weights, synthetic hidden inputs: + +| Rows | Frozen cooperative us | Shared-parameter us | Reduction | +| --- | ---: | ---: | ---: | +| 1 | 9.820 | 9.779 | 0.42% | +| 2 | 10.424 | 10.557 | -1.28% | +| 4 | 12.923 | 12.728 | 1.51% | +| 8 | 22.180 | 22.477 | -1.34% | +| 16 | 45.926 | 46.326 | -0.87% | +| 32 | 84.306 | 85.627 | -1.57% | +| 64 | 163.308 | 165.827 | -1.54% | + +The extra CTA barriers/shared accesses erase the eliminated computation in +the target batch range. This is consistent with the measurement, not an NCU +stall attribution. Register count is 123 versus 122, no spills, 1036 extra +shared bytes. Do not enable at C8/C16 or use the tiny M1 difference to replace +the established M1 route. Artifact: `gdn-shared-v1.{json,log}`. The separate +split-phase experiment also remains rejected; neither is a runtime default. + +### HC gate and repaired quality control + +The real 96-module HC graph screen stops at B4's numerical gate; no B4/8/16 +speed result is admitted. Four-rank artifacts localize the first over-envelope +output to **HC module 9, normalized input (output index 37)**. The same two +elements exceed `atol=rtol=3e-3` on every rank. The failure itself is retained +and the tolerance has not been relaxed. Artifacts: +`hc-full-b4-v1.log`, `hc-full-b4-v2.log`, and +`hc-full-b4-v2.rank{0,1,2,3}.failure.pt`. A communication-free arithmetic +replay is prepared to separate GEMM rounding propagation from IPC errors. +This numerical rejection alone does not quantify task-score regression. + +The expanded quality set was frozen before observing any expanded results: +`expanded-quality-preregistered-v1.json`, SHA256 +`72ad0a68931252ade99ac4b4ed042c8de39a33a7f175c6e085c1850efdd8903d`. +It retains the original 80 cases, registers BFCL four categories x 128, +64 schema cases, GSM8K 128, full HumanEval 164, and three fixed seed bases. +HTTP/SSE, long multi-turn fixtures and the isolated coding evaluator remain +pending. It is a registration artifact, not completed quality evidence. + +A repaired 80-case control now uses the same original prompt IDs, natural +EOS and 16K limit, with **batch HC and grouped MoE disabled** as well as the +FlashInfer parent switch. Both arms use the same adaptive #494 library. +This removes experimental components from the quality truth; do not compare +its speed to the frozen performance control. On GPU 0--3 the control scores +51/64 BFCL (15/11/14/11 by category) and 16/16 schema, zero truncations. +Candidate testing is pending; prior negative BFCL results remain on record. +Artifacts: `repaired-quality80-v1/control/`, its log and +`run-repaired-quality80.{py,sh}`. The model's normal shutdown retains one +resource-tracker shared-memory cleanup warning; this is not a leak-free +runtime admission. Model workers exited, remote services remain untouched. + +### HC arithmetic attribution and compact MoE rejection + +The communication-free HC replay now **exactly reproduces both tensors** in +the saved four-rank failure, using only the same sharded versus replicated +GEMMs and existing pointwise operations. Common-input probes first differ at +HC module 0's down projection: 696/1296 FP16 values differ, relative L2 +0.00038058. The independent chains then diverge through injection and +normalization, reaching the recorded gate failure at module 9. This isolates +that failure to projection-arithmetic association, not IPC visibility or +stale buffers. It does not rule out unrelated communication bugs or prove +model-quality noninferiority. Artifacts: `localize-hc-arithmetic.py`, +`hc-arithmetic-prefix-v1.{json,log}`. Do not relax the failed gate. + +MoE work consumption reused the existing device group table and native +W13/activation/W2/reduction math. On actual checkpoint layer-0/rank-0 weights +and captured routes, full-chain microseconds were: + +| Rows / groups | Existing | Compact W13+W2 v1 | +| --- | ---: | ---: | +| 4 / 40 | 65.338 | 76.832 | +| 8 / 71 | 97.446 | 140.448 | +| 16 / 99 | 129.696 | 164.890 | + +The first W2 mapping also colocated adjacent N tiles, confounding compaction +with the previously tested locality idea. A second screen restores the exact +original four-expert/common-N-tile CTA mapping and tests **W2 only**, without +rerunning the rejected W13 variants. It still loses: captured C16 W2 +**43.302 -> 53.523 us**, complete MoE **131.411 -> 141.158 us**, paired +within-run full-chain reduction interval **[-7.64%, -7.35%]**. All seven +changing-route graph cases preserve W13 and final output exactly, including +duplicate, invalid and empty-expert work. This rejects these fixed-grid loops, +not device planning in general; an empty-CTA count alone is insufficient +evidence of a net win. Artifacts: `moe-compact-w4-v{1,2}.{json,log}` and +the build logs/binaries. Test data are real weights/routes with synthetic +activations, not a model-quality result. + +The rejected scheduling code is retained **only as a benchmark source** in +`benchmarks/csrc/sm70_moe_compact_tasks.cu`. The existing production +`nvfp4_grouped_decode_sm70.cu` was restored unchanged. No CMake or serving +dispatch selects the benchmark namespace. The pre-retirement source delta is +retained in `moe-compact-before-retirement.patch` for exact binary provenance. + +### Repaired original-80 combined result: still not admitted + +The candidate runs on the subsequently freed GPU 4--7; the control used GPU +0--3, same V100 host, weights, sampling, prompt IDs, cache/state contract and +fixed cohorts. This is a quality localization run, **not a same-GPU speed +comparison**. All four candidate workers select GDN, MQA and sparse QSA; +batch HC and grouped MoE remain disabled in both arms. + +| Original subset | Production control | Combined candidate | +| --- | ---: | ---: | +| BFCL simple | 15/16 | 15/16 | +| BFCL parallel | 11/16 | 10/16 | +| BFCL multiple | 14/16 | 14/16 | +| BFCL irrelevance | 11/16 | 11/16 | +| JSON Schema | 16/16 | 16/16 | + +No truncations; all 80 first tokens match. Twelve complete outputs differ. +Only `parallel_1` changes pass/fail: the candidate emits one valid tool call +instead of the required two. Total BFCL is **51/64 -> 50/64**. The paired +score delta is -1.5625 percentage points; a conservative 95% interval from +simultaneous exact-binomial improvement/regression bounds is +**[-9.5612, +6.5981] pp**. The interval is reported, not used to dismiss the +negative result or declare noninferiority. This original-80 offline screen +is not the registered expanded or HTTP/SSE suite. + +Artifacts: `repaired-quality80-v1/{control,candidate}/`, both logs, +`comparison-v1.{json,log}`, and `compare-repaired-quality80.py`. The original +adverse results remain intact. A scorer-only ablation uses the same fixed +80 cases and no new seed selection to distinguish MQA from GDN/sparse-QSA +effects. Full wheel, extended quality and endpoint targets remain unmet; +keep #523 Draft and `VLLM_SM70_FLASHINFER_BATCH=0` by default. + +### Original-80 ablations and packaged GDN follow-up + +The same repaired control and original cases now have three additional +ablations. No seed, prompt, output limit, or failed case was changed: + +| Arm | BFCL / 64 | Schema / 16 | New BFCL failures vs control | +| --- | ---: | ---: | --- | +| Production control | 51 | 16 | n/a | +| Device-planned MQA only | 51 | 16 | none | +| GDN only, prototype library | 52 | 16 | none | +| MQA + wheel-component GDN, no native sparse QSA | 52 | 16 | none | +| MQA + GDN + native sparse QSA | 50 | 16 | `parallel_1` | + +MQA alone preserves **all 80 complete output token lists**, not just scores. +GDN-only and MQA+GDN improve `irrelevance_4` and have no newly failing cases; +their 80 complete output token lists also match each other exactly. +Their paired BFCL difference is +1.5625 pp, conservative 95% interval +**[-6.5981, +9.5612] pp**. All arms have zero truncations and identical first +tokens. These small offline ablations narrow the adverse signal to native +sparse QSA and its interaction; they do not admit GDN, prove broad +noninferiority, or substitute for registered multi-seed/HTTP/SSE tests. +Artifacts: `repaired-quality80-v1/{mqa,gdn,mqa_gdn}/` and +`comparison-v3.{json,log}`. Earlier comparisons remain intact. + +The sparse arithmetic audit identifies a concrete contract difference to +investigate next. Production Triton attention rounds softmax probabilities +to FP16 before its PV dot while retaining an FP32 denominator; the pinned +FlashInfer decoder accumulates FP32 probabilities against converted values. +Its tile/split reduction order also differs. This is a source-level lead, +**not causal proof that one cast caused the missing tool call**. Blindly +adding that cast does not reproduce the other reduction boundaries; a more +accurate FP64 operator comparison does not clear the observed model failure. + +GDN now has a formal `_sm70_flashinfer_gdn_C` CMake/setup component, with +separate C++ namespaces for H2560/Q4/V12, Q8/V24 and Q16/V48. The benchmark +and package reuse one binding/header rather than diverging implementations. +The loader resolves the installed component before capture; explicit +preloaded prototypes prevent duplicate registration, and missing geometry +falls back locally. No external path or worker-side compilation is required +for this component. The source attribution now names the exact upstream +experimental GDN path. Other shapes remain unsupported, not globally gated +by TP degree or model name. + +The staged CMake component was loaded through normal package discovery in +the MQA+GDN model arm above, without the external GDN-library override. +Its SHA256 is +`4003e0393021456924606ebfe93eb62f4d1dfbbfabedbccb373d859b170295c1`. +Seven GPU graph tests pass for B1/2/4/8/16/32/64: all three head partitions +give bitwise identical outputs, conv states and FP32 recurrent states across +eight changing-input steps, including slot permutation, negative padding, +SD conv layout and non-dense pool strides. This tests operator geometry +isolation, **not distributed TP communication**. Artifact: +`native-gdn-gates-v1.log`; build/install/import provenance is retained in +`native-gdn-*-v1.log` and `native-gdn-stage/`. + +The final formatted source rebuild is staged separately in +`native-gdn-stage-v2/`, SHA256 +`7cae0a12e1018c44e7a9d5f62a7feec85cd85bbeba92c805d18aa72c2cf2988c`. +It imports all three namespaces without initializing CUDA. All 30,537 lines +of disassembled SASS match the tested v1 component exactly; debug/source +metadata changes the library hash. Artifacts: `native-gdn-build-v2.log`, +`native-gdn-install-v2.log`, `native-gdn-v{1,2}-real.sass`. The initial +disassembler lookup under the build CUDA shim failed because that shim has +no `cuobjdump`; the successful comparison uses `/usr/bin/cuobjdump`, not +the empty outputs of that failed lookup. Do not relabel the v1 model run as +a new v2 wheel-install test. + +A reload audit also found that an already prepared layer could retain its +old derived BA weights if a later reload changed to an unsupported dtype, +layout or geometry. Initial unsupported layers still fall back, but changing +an already captured contract now fails explicitly and requires graph +rebuilding; silently leaving `_sm70_fi_ready` with stale weights is forbidden. +The CPU routing/statistics/compiled-boundary suite now passes **42 tests** +(`fi-cpu-v6.log`), including absent/preloaded component and reload guards. + +Coverage status for this follow-up: + +| Hot chain | Implemented / evidenced | Remaining gate | +| --- | --- | --- | +| QSA scoring/selection | Device work plan; exact production top-k/index chain wins | Wider routing, long-context and expanded quality | +| QSA prefill plan | Reused #494; adaptive stable sort recovers overhead | Wider prefill timing and racecheck | +| Sparse attention | Native prototype reaches actual graph execution | Retained BFCL failure; arithmetic/interaction isolation; packaging | +| GDN | Native component; projection/BA/conv/state boundary; staged model and geometry tests | Expanded quality, full-chain B-shape admission, full wheel | +| HC / TP | Full-chain B4 numerical failure localized to projection association | Preserve arithmetic before further overlap work; multistream/IPC gates | +| MoE | Existing grouping retained; two compact-loop designs measured and rejected | New critical-path evidence before another scheduler change | + +The split-phase/shared-parameter GDN variants and compact MoE loops are not +production selections. The full installable wheel, clean-environment +acceptance, expanded quality, long-context capacity and new fixed-contract +end-to-end targets remain unfinished. No 20--30% model-throughput gain or +release default change is claimed. Keep the Draft scope and parent default +off; do not alter the remote deployment. + +Publishing preflight: `onecat/main` advanced to +`4366d9d5fe80eeaf79575b51ec36a6a032673df0`. The measured branch is not +silently rebased onto that changing integration line; merge compatibility +and a fresh integration gate remain required before promotion. Applicable +source checks pass (`hot-chain-precommit-v8.log`). All model, benchmark and +waiter processes owned by this follow-up have exited. Subsequent GPU 0--3 +and 4--7 allocations belong to other task leases and were not interrupted; +there is no task-owned resident API or new remote service. + +### Numerical isolation follow-up: EOS boundary and HC projection stages + +The last measured E2E comparison remains C4/C8/C16 +**232.169/390.046/630.996 tok/s**, +6.98%/+7.00%/+7.35% versus the original +performance control. It is still not quality-admitted. The newer MQA and +stable-planner savings have **not** been measured in a new E2E campaign. + +Inspecting the retained repaired-quality outputs further localizes +`parallel_1`: MQA+GDN and the sparse-QSA combined arm share their first +**51 output tokens** exactly. The first tool call is complete and identical. +At zero-based offset 51, MQA+GDN emits newline token `198`, then a second +`calculate_em_force` call with `d_time=10`; native sparse QSA emits EOS +`248046`. Completion lengths are 105 versus 52 tokens. This is premature +generation termination, **not a parser dropping a generated second call**. +It does not identify which numerical change shifted the EOS draw. + +The new CPU counterfactual retains Q/K/V, index order, duplicate/invalid +selections, reference 16-column split boundaries and the output/gate casts. +Everything else is evaluated in FP64, changing only the probability +materialization before PV. On all **36** previously retained QSA inputs, +restoring the FP16 probability cast is closer to the recorded production +output: median relative L2 **1.27622e-4 -> 3.82219e-5**, median per-case error +reduction **70.34%**. This supports that rounding as a substantial source of +the *operator* difference. These inputs came from the earlier shadow run, +not the repaired `parallel_1` EOS step; CPU arithmetic is not an exact CUDA +emulator. Do not interpret the result as permission to blindly round the +native kernel, reduce state precision, or declare model noninferiority. + +Reproducible numerical tool and sanity tests: + +```bash +CUDA_VISIBLE_DEVICES='' .venv/bin/python \ + benchmarks/kernels/benchmark_sm70_qsa_rounding_isolation.py \ + --captures .artifacts/quality-shadow-v2 \ + --out .artifacts/qsa-rounding-isolation-v1.json +CUDA_VISIBLE_DEVICES='' .venv/bin/python -m pytest -q \ + tests/kernels/test_sm70_qsa_rounding_isolation.py \ + flashinfer-sm70/tests/test_batch_routing.py \ + flashinfer-sm70/tests/test_compiled_gdn_boundary.py \ + flashinfer-sm70/tests/test_paired_stats.py +``` + +The combined CPU suite passes **45 tests** (`arithmetic-routing-cpu-v1.log`), +including three new tests for duplicate weighting, invalid/NaN padding and +keeping the denominator unrounded. The counterfactual artifact is +`qsa-rounding-isolation-v1.{json,log}`, JSON SHA256 +`05d6b8e1704feaf39503cbbfdbb52bcef673098719d428ffd46679c2546d1db9`. + +For HC, `benchmark_sm70_hc_arithmetic_isolation.py` now separates full +sharding, down-only sharding, and up-only sharding across the 96-module +arithmetic chain on one GPU. It must first reproduce the saved four-rank +failure exactly and retains the same 3e-3 absolute/relative envelope. +It does not measure collective performance or replace the distributed/full +model gate. Retaining the reference down projection while sharding only up +is a test hypothesis, **not an implemented or validated serving repair**. + +A focused, diagnostic-only teacher manifest also freezes the original +16-case parallel cohort through offset 51, using the existing per-case seeds +and trajectories. It will compare raw/processed logits and the stateless +EOS-versus-newline draw on all four ranks, with A/A repeats, before another +natural quality run. Manifest `parallel-eos-teacher-manifest-v1.json`, SHA256 +`2f179c9940d28bd9a631821bd669926e5b11d5f29423f6f2a47f2092f9be01d2`. +Launchers/hooks stay in task artifacts and require explicit diagnostic +environment variables; they are not serving code. The spawn import check +passes without CUDA initialization. **This model diagnostic has not run.** + +The HC launcher waited its ten-minute lease window and returned code 75; +a fresh attempt was also blocked by live foreign GPU processes, even after +one lease owner exited. No HC GPU result file was produced and no model +diagnostic was launched. The waiter has exited; no task-owned GPU process +or resident service remains. Do not preempt other tasks or describe these +queued checks as passed. Resume HC single-GPU isolation and then the focused +four-rank EOS diagnostic after resources are actually free. Source runtime +defaults remain unchanged; fixed quality and E2E admission are still pending. + +### Implemented numerical repairs and new micro gates + +After the preceding lease-only attempt, resources became available. HC stage +isolation now reproduces the saved failure and completes 16 inputs across all +96 real HC weights. Up-only sharding is bitwise identical at every arithmetic +intermediate; down-only and combined sharding reproduce the drift. The runtime +now preserves the original full `336 x 10240` down GEMM and its FP16 output, +then shards only up and retains the existing fused mix/gather. It no longer +runs the down-shard collective. This is not an IPC workaround or precision +change. M1 and prefill fallbacks are untouched. + +Four-rank complete HC Graph replays pass the existing **every-intermediate +atol=rtol=3e-3** gate at B4/8/16, with 16 changing inputs plus post-timing +checks. Unlike the single-device projection isolation, the full native chain +is not bitwise identical. Five alternating paired timing blocks give: + +| Rows | Original HC chain ms | Repaired chain ms | Paired reduction, 95% CI | +|---|---:|---:|---:| +| 4 | 3.68872 | 3.67303 | 0.422% [0.407%, 0.438%] | +| 8 | 3.79545 | 3.78903 | 0.213% [0.154%, 0.272%] | +| 16 | 4.01596 | 4.01972 | -0.092% [-0.116%, -0.068%] | + +These include all HC norms and the final mixer, not attention/MoE/PLE +computation. B16 has **no speed admission**. HC remains opt-in; do not claim +that preserving correctness also preserved the old sharded-down speedup. +Artifacts: `hc-projection-ablation-v1` and `hc-up-only-b{4,8,16}-v1`. + +The native sparse-QSA compatibility experiment uses Volta WMMA FP32 QK/PV +accumulation, the production 16-token tile partition and an explicit FP16 P +boundary, retaining the FP32 denominator, ordered selections and original +output gate. Pinned FlashInfer virtual-page preparation and FP32 cascade +remain in use. This is a separate native namespace, not a Triton call hidden +under a FlashInfer name. No EOS suppression or sampling change is introduced. + +On 36 retained real QSA inputs, median relative L2 versus production changes +from **1.27593e-4 to 5.34201e-5**. Eight changing Graph replays at each of +B1/2/4/8/16/32/64 pass the 2e-3 envelope, including empty/duplicate/invalid +selections and relocated pages. Whole-QSA micro timings including preparation +and merge improve over production by B4 **6.62%**, B8 **38.27%**, B16 **49.57%**; +paired 95% CIs are [6.38%,6.86%], [37.93%,38.59%], [49.43%,49.71%]. The new +kernel is slower than the numerically different SIMT experiment; that speed +tradeoff is recorded, not hidden. Artifact: `qsa-compat-replay-v1`. + +Runtime capability dispatch now recognizes the compatibility namespace and +uses its matching split partition and call-local scratch. Explicit old SIMT +preloads remain available for the registered counterfactual. CPU routing and +HC tests: **84 passed**. Actual compatibility bridge with output gating and +changing-input CUDA Graph: **2 passed**. These checks do not establish model +quality, repair the EOS case by themselves, or constitute new E2E numbers. +The focused four-rank diagnostic and natural 80-case model gates are running; +the experimental parent remains default off and #523 remains Draft. + +The dynamic-shape audit then found that the first compatibility wrapper's +fixed B4/B8/B16 split rule did not cover intermediate live widths, multiple +KV heads, or short selections. The `repaired-quality80-v2/compat_hc` model +attempt was stopped before producing quality/speed results, and its artifacts +are retained. Reuse `_qsa_sparse_launch_profile` and its maximum-useful-split +rule instead. Expanded CPU checks: **118 passed**; real bridge tests at +B4/8/9/15/16 with one/two KV heads and selection lengths 15/65/2051: +**20 passed**, including eager/Graph equality and the existing output gate. +Do not mistake the initial fixed-width micro timing for dynamic-shape coverage. + +Both registered EOS teacher-prefix arms now completed twice. All four ranks +agree bitwise, both A/A repeats agree, and the arms' first-token raw logits +are bitwise identical. At offset 51 (position 468, seed 20260922), the +production-sparse MQA+GDN arm's [newline, EOS] raw logits are +**[26.09375, 27.25]**, versus **[26.015625, 27.671875]** for old SIMT QSA. +After the unchanged top-k/top-p processing, probabilities are respectively +**[0.239349, 0.760651]** and **[0.160266, 0.839734]**. The unchanged stateless +draw selects newline in the first arm and EOS in the second, in every rank +and repeat. This proves a decode-induced probability shift, not a parser, +seed or first-prefill divergence. It does not isolate every arithmetic +suboperation or prove the replacement passes natural quality. Captures and +summary: `parallel-eos-v1/{mqa_gdn,candidate,summary-v1.json}`. + +The fresh combined natural run is `repaired-quality80-v3/compat_hc`; it retains +original prompts, seeds and max16K/natural EOS, with grouped-MoE experiments +off to isolate these repairs. Only after each original suite meets the +repaired production control and the EOS case passes does its driver reuse +the same engine for the unchanged 8K/256 fixed-width speed runner (SHA256 +`c969b3e8087c2777fc6857c2b30c6aa81ed54558cef6c64c181c2ca1f4a60c9a`). +This is an original80 prerequisite, **not** expanded multi-seed quality +admission. The native QSA experiment binary used in these checks has SHA256 +`a0b12d034bc12569c85ecfda68369b6ca26b700929627c37ca97f49bafa7be17`. diff --git a/docs/design/sm70_flashinfer_layer_fusion.md b/docs/design/sm70_flashinfer_layer_fusion.md new file mode 100644 index 0000000000..234c426e96 --- /dev/null +++ b/docs/design/sm70_flashinfer_layer_fusion.md @@ -0,0 +1,197 @@ +# FlashInfer-derived GDN / HC layer-fusion campaign + +## Purpose + +Finish operator adaptation and evidence first, then run one consolidated +end-to-end comparison. Do not launch another full model per operator. Maintain +the same no-MTP FlashNext workload, output-quality gates and fixed-70 +concurrency targets (C4/C8/C16 aggregate 238/420/728 tok/s). + +Integration base: `onecat/main` at `755baae1d075ee04fa9096b23fc0225b23589a86`. +Owned branch/worktree: `codex/v100-flashinfer-gdn-conv-20260905-173007` / +`worktrees/v100-flashinfer-gdn-conv-20260905-173007`. +Related QSA prototype is Draft #513; prior batch HC work is Draft #504. +This change does not duplicate their QSA or TP-sharding implementations. +The M1-only #506/#510 work remains separate and is not overwritten. + +Source pin: FlashInfer `6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f`. +The CUDA code is genuinely derived from its kernels, not a FlashInfer API +redirect to Triton. The adapters are benchmark-only until gates pass. + +## GDN adaptation + +Derived from `gdn_kernels/experimental/kernel/gdn_fused_decode_sm120.cu`: +gate B/A projection, width-4 causal convolution, Q/K normalization, gating, +FP32 delta-rule state update and attention output in one kernel. FP16 +activations/weights and FP32 recurrent state remain unchanged. Geometry is +JIT-specialized, not identified by checkpoint name or global server settings. + +The actual FlashNext TP4 shard is hidden=2560, Hq=4, Hv=12, D=128, +QKV-width=2560 and BA-width=24. Do not reuse an old Hq4/Hv8 GDN result as this +model's baseline. Reference QLA source SHA256: +`fd6389cef9f1b38df7e122e582221d74d9ae1fba377ac3bec0da047fd3d30af8`. + +Load-bearing integration changes: + +- Production conv PTX emits `mul.f16`, then `cvt.f32.f16`. Preserve the FP16 + product boundary and ordered FP32 accumulation, not upstream BF16/FP32 + widened multiplication. Keep the SiLU-to-FP16 materialization. +- Preserve B/A projection's FP16 materialization before FP32 sigmoid/softplus. +- Preserve in-place state aliasing and V-major `[pool,Hv,V,K]` layout, + explicit pool strides, DS/SD conv layouts and strided QKV projection views. +- Negative padding owns no state and emits zero. Valid slots must be unique + within a call, in range, and owned by the scheduler; this prototype is not + a replacement for prefix-cache copy-on-write / metadata validation. +- Use cooperative launch and CUDA grid synchronization instead of relying on + a software spin barrier's regular-launch residency assumption in a runtime + with auxiliary streams. Cooperative graphs/stream capture are supported + by CUDA; see [NVIDIA's CUDA 11 description](https://developer.nvidia.com/blog/cuda-11-features-revealed/). + Runtime capability and exact occupancy are still checked before launch. + +## HC adaptation + +The fused combine/Gemma-norm adapter derives from FlashInfer's +`FusedAddRMSNormKernel` vector IO, shared staging and warp/block reduction. +It adds HC's injection gate and per-branch/shared affine addressing. The +materialized residual is rounded before RMS statistics, and Gemma affine is +`fma(y, w, y)`, not an unqualified ordinary RMSNorm replacement. Single-lane +shared reduction writes avoid redundant same-address writes. + +FP16/FP32 residual and block types are distinct template parameters. Weight +and injection dtypes match; no FP8/int8/QPN approximation is introduced. +The initial component accepts vector-aligned contiguous matrices with +group width a multiple of 8, <=4096, and 1/2/4/8 warps for screening. No +production settings or server max-seqs/TP/prefill limits are changed. + +## Existing negative evidence not to repeat + +- Previous v0.6.13 standalone GDN-input GEMV replacement reached parity, + and M1 HC down/up replacement regressed. A route hit is not a gain. +- Batch GDN projection concatenation/overlap and generic cuBLASLt retuning + already failed their integration thresholds. +- FP16 recurrent-state compression caused large state/output error in the + prior batch audit and is not included. +- HC projections plus communication already have an isolated #504 candidate; + reuse its validated pieces instead of reimplementing or counting them twice. + +## Test plan and progress + +Environment: Python 3.12.13, Torch 2.10.0+cu128, CUDA 12.8, native SM70, +private Torch/Triton caches. No service/model engine is launched at this stage. + +1. Native CPU compilation and load, current-library/geometry provenance. +2. Independent GDN oracle; dynamic state slots, live slot zero, padding, + poisoned outputs, strided input/state layout and graph/eager equality. +3. Checkpoint-weight GDN component against actual production conv+FlashQLA; + separate conv/state/output errors and preserve reference trajectories. + An identical input-refresh copy is timed in both arms because production + conv mutates its input; QKV/Z projection itself is excluded. +4. HC full combine/norm versus current Triton, FP16/FP32 residuals, B1/4/8/16. +5. Memcheck/racecheck/synccheck and longer independent recurrent histories + before runtime admission. Micro error thresholds are screening gates, + not proof of task-level quality non-inferiority. +6. Integrate only winners together with QSA, preserve fallback routes and + baseline precision. Then one consolidated E2E C1/4/8/16 comparison plus + coding/tool/schema quality checks, reporting actual routes and pure decode + separately from prefill/TTFT. Failed operators stay off. + +Native GDN and HC SM70 builds passed. Cooperative GDN builds also passed, +including Hq4/Hv12 and Hq8/Hv24 modules loaded in the same process; +per-geometry Torch namespaces prevent duplicate registration. + +### GPU results, 2026-09-06 + +After the prior lease released GPU 0--3, the `72af224161` tests ran on locked +GPU 0: **16 passed**, including all 8 GPU cases (33.39 s, including HC build). +This covers independent GDN histories, DS/SD conv layouts, strided QKV, +padding/live slot zero, and HC shared/per-branch weights, mixed FP16/FP32, +non-power-of-two widths, poisoned outputs and graph replay. +Artifact: `.artifacts/gdn-hc-unit-gpu-v4.log`. + +Checkpoint component screen: RadixArk Qwen3.8-Flash-Next-NVFP4, layer 0, +TP4 rank-0-shaped **FP16** GDN weights, FP32 state, Hq4/Hv12/D128, +synthetic changing hidden states. The target MoE quantization is not being +retested here. CUDA Graph, 9 alternating paired samples, 30 calls per graph, +identical raw-QKV refresh in both arms; no full model launch. GPU clocks use +automatic boost and were not locked, so use paired deltas, not cross-run +absolute comparisons. First screen uses 16 local steps and then 256 fully +independent state-history steps with padding and slot recycling. + +| Rows | Existing BA + conv + FlashQLA, us | Fused chain, us | Latency reduction | +| --- | ---: | ---: | ---: | +| 1 | 17.237 | 11.674 | 32.28% | +| 4 | 28.604 | 15.087 | 47.26% | +| 8 | 33.041 | 22.801 | 30.99% | +| 16 | 57.344 | 47.343 | 17.44% | + +All four independent-history screens pass the unchanged operator gates. +Worst output relative L2 across those histories is 2.458e-4 and worst FP32 +state relative L2 is 2.636e-5; conv state updates remain exact. These errors +are not zero and this is not a task-quality admission. The M1 reference +includes separate BA projection whereas production M1 can fuse QKVZ/BA; +**do not count the M1 component delta as a production gain**. +Artifact: `.artifacts/gdn-screen-v1.log`. + +Commands inside the owned GPU-lock/environment launcher (the wrapper sets +the pinned QLA binary, FlashInfer headers and task-private caches): + +```bash +.venv/bin/python -m pytest -q -x --confcutdir=flashinfer-sm70/tests \ + flashinfer-sm70/tests/test_layer_fusion.py +.venv/bin/python -m benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv \ + --model /path/to/Qwen3.8-Flash-Next-NVFP4 --steps 16 +.venv/bin/python -m benchmarks.kernels.benchmark_sm70_flashinfer_hc_norm +``` + +Reference QLA binary SHA256: +`3982305151798be22a1dabd0140feb085f787e1e76da01f58d8256de66050975`. +GDN candidate binary SHA256: +`45cfe0090f43792ac8f4a5a21b9475e988c0f6872c874faf2e5da967db978d66`. +The 8-row-warp cooperative kernel has 122 registers/thread (M1: 110), +zero stack/local memory. These are static resources, not measured occupancy. + +HC shared-staging version passed all numerical screens but **lost every +timed shape**. FP16 residual results (best of 1/2/4/8 warps): + +| Rows | Existing Triton, us | Best FlashInfer-derived HC, us | +| --- | ---: | ---: | +| 1 | 3.164 | 4.321 | +| 4 | 3.052 | 3.942 | +| 8 | 2.918 | 3.717 | +| 16 | 3.144 | 3.953 | + +The FP32-residual arm also regressed. Retain the existing production HC; +do not promote this version on the basis of a FlashInfer label. +Artifact: `.artifacts/hc-norm-screen-v1.log`. HC v1 binary SHA256: +`ef905abb8feb41a5887fc64dc45f11dbea97039c60095a1e17b93fbe143d079b`. + +### Follow-up candidates and remaining gate + +- GDN four-row-per-warp variant: B8/16 independent-history screens pass; + paired baseline/candidate medians 36.420/25.054 and 57.344/46.353 us. + Registers fall to 89 (M1: 78), still no spills. This is not a same-run + eight-versus-four comparison; do not change the default based on the small + cross-run difference. Artifact: `.artifacts/gdn-screen-r4-v1.log`. +- HC register-staging variant: retain materialized FP16/FP32 residuals in + registers over the reduction, avoiding the shared-value write/reload. + Local D2560, 4/8-warp specialization; all other geometries retain the + general component. Native build passes, but GPU correctness/speed is pending. + Artifact: `.artifacts/hc-norm-build-v2.log`; binary SHA256: + `6087639c2f615ce04775000d556325589f391993c88653a5e3baad04bff347ff`. +- Updated CPU suite: **8 passed, 12 GPU cases skipped** with GPU hidden. + The four additional register-HC cases have not run on GPU. Previous v1 GPU + results must not be relabeled as v2 validation. +- The targeted GDN memcheck attempt did **not** launch: the new QUASAR E4M3 + task acquired the paper GPU 0--3 lease between component jobs. Exit 75 and + an empty sanitizer log are not a pass. Wait for release; never preempt it. +- No runtime integration, new E2E throughput, or model-quality pass yet. + Preserve existing projection/M1 and HC paths until complete-chain gates + pass; avoid recomputing BA if using the new fused-input GDN boundary. + +Local artifacts: `.artifacts/gdn-build-v1.log`, `gdn-build-v2.log`, +`gdn-build-cooperative.log`, `gdn-build-multi-geometry.log`, +`hc-norm-build-v1.log`. GPU test logs may exist +but be empty when lock acquisition timed out; file presence is not a result. +No model speed or output-quality result is claimed yet. No owned service. + +AI-assisted work (Codex); human review and DCO sign-off required before merge. diff --git a/docs/design/sm70_flashinfer_qsa_port.md b/docs/design/sm70_flashinfer_qsa_port.md new file mode 100644 index 0000000000..9f797d77d7 --- /dev/null +++ b/docs/design/sm70_flashinfer_qsa_port.md @@ -0,0 +1,207 @@ +# FlashInfer CUDA QSA port to Volta + +## Purpose and frozen baseline + +Port actual FlashInfer computation to SM70, starting with sparse QSA decode. +Changing a backend name or forwarding to the existing Triton/Flash-V100 kernel +does not meet this objective. This first PR is an isolated benchmark prototype; +it does not select a new serving backend or change release defaults. + +- Integration: `1CatAI/1Cat-vLLM`, `onecat/main`. +- Base: `755baae1d075ee04fa9096b23fc0225b23589a86`. +- Branch: `codex/v100-flashinfer-qsa-sm70-20260905-164136`. +- Draft PR: [#513](https://github.com/1CatAI/1Cat-vLLM/pull/513), initial + implementation commit `6b2f4ad8e9c1ddebf375bb1ea163dc2b5891ced4`. +- FlashInfer source: [6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f](https://github.com/flashinfer-ai/flashinfer/tree/6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f). +- Its CCCL submodule: `16bd510c9b712e82b0ab6cbb630d8e29ba1f7116`. +- The previous WMMA primitive probe has its own older source pin; it is not + silently upgraded or used as evidence for this new attention path. + +Scope is different from the HC/QSA follow-up #504 and M1 QSA #507: this +instantiates FlashInfer's CUDA kernels, not another existing native-FA variant. +No MoE/HC changes from those worktrees are copied here. + +The fixed no-MTP concurrency goals remain 238/420/728 aggregate decode tok/s at +C4/C8/C16, respectively (85/75/65% of C times 70 tok/s). The previously recorded +C16 model step was 27.125 ms; meeting 728 tok/s requires 21.978 ms. Component +microseconds must not be presented as a new model result. + +## Actual implementation + +`flashinfer-sm70/include/flashinfer/attention/sm70/qsa_decode.cuh` provides a +virtual page-size-one sparse cache adapter. It directly instantiates the pinned +upstream `BatchDecodeWithPagedKVCacheKernel` and `MergeStatesKernel` from +`include/flashinfer/attention/decode.cuh` and `cascade.cuh`. The upstream header +tree is unmodified. QK, online softmax, PV, split-state normalization and merge +therefore execute FlashInfer's actual CUDA implementation. + +Three GPU operations are measured together: + +1. Expand raw logical QSA selections into physical 64-bit offsets and prepare + persistent split metadata. Preserve order and duplicate weighting. +2. Execute FlashInfer paged decode using these sparse virtual pages. +3. Merge FP32 partial outputs/LSE with FlashInfer's cascade kernel; cast once + to FP16 at the final output. + +SM70 compatibility is obtained by instantiating the upstream SIMT/GQA kernel. +Its `cp_async.cuh` already has ordinary vector-load/store fallback below SM80, +and the decode body retains block barriers. This first candidate does not use +Tensor Cores and does not claim to replace tensor-core instructions by equally +fast software emulation. If SIMT loses, optimize tile reuse and evaluate native +Volta WMMA with the same softmax/masking contract before considering dispatch. + +The Torch extension needs a consistent CCCL include set (Thrust, CUB and +libcudacxx) and undefines Torch's disabled CUDA half-conversion/operator macros. +These are build integration fixes, not a global relaxation of FlashInfer's +supported-GPU checks. No FlashInfer Python package is installed or shadowed. + +The prototype covers FP16 Q/K/V, D256, 1..32 query heads, GQA groups 1/2/4/6/8, +1..64 splits, variable batch, selection width, page size and request mapping. +It accepts vector-aligned strided tensors and rejects misalignment. This is a +component contract, not a global max-seqs/TP/chunk/prefix-cache restriction. +Other dtypes, fused output gate, indexer, prefill and serving integration are +not implemented by this PR's first iteration. + +QSA causality is already encoded by the existing selector/expander. Do not +reinterpret a query row as dense causal attention or physically sort pages. +Negative/out-of-range requests, indices, pages and split padding are masked. +Invalid cache loads use a persistent zero vector, not an arbitrary cache page: +zero probability multiplied by a cached NaN would still produce NaN. + +## Test plan + +No parameter change, routing-policy change, quantization or model approximation +is allowed as a substitute for a correct faster operator. + +- Compare against an independent FP32 oracle that preserves repeated indices. +- Test empty/invalid rows and pages, tail and split padding, GQA layouts, + non-contiguous tensors, and input/output pointer alignment. +- Poison output/partial/metadata buffers, mutate indices and physical mappings, + then replay a captured graph. Require replay to equal this implementation's + eager result and both to match the FP32 oracle (`atol=2e-3, rtol=1e-2`). +- For timed model-shaped cases require relative L2 error <= 5e-3. This is an + operator screen, not a substitute for task-level quality tests. +- Time current two-warp Triton and FlashInfer with identical inputs, including + preparation/merge. Record source SHA256, raw paired samples, B1/4/8/16, + CUDA Graph and SM70. Reject results if another process uses the GPU. +- Only after a stable component gain: sanitizer, runtime metadata integration, + representative C1/4/8/16 E2E and output-quality regression. No default change + before these gates; no GPU-serving process left resident after testing. + +## Reproduction + +Use the task virtual environment backed by Torch 2.10.0+cu128, Python 3.12 and +CUDA 12.8, not system Python. Do not use CUDA 13 to compile SM70. + +```bash +bash tools/prepare-flashinfer-sm70-qsa.sh +export CUDA_HOME=/path/to/cuda-12.8 +export TORCH_CUDA_ARCH_LIST=7.0 +export TORCH_EXTENSIONS_DIR="$PWD/.artifacts/torch-extensions" +export MAX_JOBS=2 +# CPU-only compilation; no GPU or model loaded. +CUDA_VISIBLE_DEVICES='' .venv/bin/python -m benchmarks.kernels.flashinfer_sm70_qsa +CUDA_VISIBLE_DEVICES='' .venv/bin/python -m pytest -q \ + --confcutdir=flashinfer-sm70/tests flashinfer-sm70/tests/test_qsa_decode.py +# Acquire the shared GPU ownership locks and confirm an idle physical GPU first. +CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest -q \ + --confcutdir=flashinfer-sm70/tests flashinfer-sm70/tests/test_qsa_decode.py +CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m \ + benchmarks.kernels.benchmark_sm70_flashinfer_qsa --compare-triton +``` + +The comparison needs the source tree's usual working vLLM environment; the +FlashInfer-only tests require neither vLLM's native extension nor Flash-V100. + +## Test result / status + +- CPU compilation and shared-library loading: passed for native `sm_70`. +- `cuobjdump` confirms native `sm_70`; GQA6 decode uses 72 registers/thread + and reports zero stack/local memory (not a measured occupancy or speed gain). +- CPU tests: **5 passed, 9 GPU tests skipped** (not counted as GPU passes). +- Subsequently acquired an idle, locked V100: **14 tests passed**, including + all 9 GPU cases. Runtime: 2.82 seconds; graph replay equals eager. +- Python Ruff check/format and targeted C++/CUDA formatting: passed. +- All applicable staged pre-commit hooks: passed, including mypy, shellcheck + and project-specific checks. +- Compute Sanitizer 12.8.93 memcheck: 2 selected GQA6 cases passed, **0 errors**. + Racecheck of the same cases: **0 errors, 0 warnings**. Both include mutated + metadata, graph replay and empty rows. This is targeted, not whole-model + sanitizer coverage. +- Serving dispatch/defaults/quality or model speed: unchanged, not validated + by this prototype. Do not advertise a throughput increase yet. + +Build and test outputs are stored in this task's unversioned `.artifacts/` directory +(`flashinfer-qsa-build-v4.log`, `flashinfer-qsa-cpu-tests-v1.log`, +`flashinfer-qsa-resources-v1.txt`, `flashinfer-qsa-tests-v1.log`, +`flashinfer-qsa-memcheck-v3.log`, `flashinfer-qsa-racecheck-v1.log`). The initial +build failures from mixed CCCL headers and disabled half conversions were +localized and fixed; do not rerun those variants as performance experiments. + +The system's 2022.4.1 sanitizer could not instrument this runtime (missing +injection-library lookup, then application-exit failure). Neither failure is +counted as a pass. The successful run uses NVIDIA's CUDA 12.8.1 redistribution, +`cuda_sanitizer_api` 12.8.93, archive SHA256 +`ae3574f052c0e06c95305962668eb1fe6ab571dfbb58b305fdb14d523bb1b240`, unpacked +only in task `.deps/`. It does not replace system tools. Command: + +```bash +compute-sanitizer --tool memcheck --kernel-name kns=flashinfer \ + --error-exitcode 99 .venv/bin/python -m pytest -q \ + --confcutdir=flashinfer-sm70/tests flashinfer-sm70/tests/test_qsa_decode.py \ + -k '16-2051 or 4-33' +# Repeat with --tool racecheck, using the same pinned tool. +``` + +### First paired component screen + +After graph and GPU-clock warmup, 11 alternating-order paired samples of 100 +calls per graph, independent 8192-token requests, 2051 selection slots, page +size 784, FP16 Q/K/V, local Hq=6/Hkv=1/D256. GPU: V100-SXM2-32GB; driver +580.173.02; post-sampling SM/memory clocks 1530/877 MHz for each row count, +automatic boost (not clock-locked); temperature 36..40 C. This is a component +test on one TP4-shaped shard, not a TP4 model timing or a 1200 MHz comparison. + +| Query rows | Current two-warp Triton, us | Best FlashInfer, us | Splits | Latency change | +| --- | ---: | ---: | ---: | ---: | +| 1 | 27.228 | 30.556 | 32 | **+12.22% regression** | +| 4 | 54.641 | 45.066 | 32 | -17.52% | +| 8 | 100.055 | 60.611 | 32 | -39.42% | +| 16 | 167.311 | 87.491 | 16 | -47.71% | + +Every FlashInfer timing includes preparation, actual upstream decode and merge. +The split counts are selected from an explicitly exploratory 16/32/64 sweep, +not a validated production policy. First-run cold-clock numbers are retained +in `flashinfer-qsa-screen-v1.log`, but the warmed screen above is +`flashinfer-qsa-screen-v2.log`. Do not compare unrelated historical absolute +microseconds. The reference QSA source SHA256 is +`01e9e97d49fe750e5e0d2ee61961e53349ba631a67be0df22e007a715c625543`. + +All four query-scale/tail/mapping cycles passed in every timed shape and split +choice; maximum FlashInfer relative L2 error against FP32 was 0.00021144. +This does not establish token equality or model-quality non-inferiority. + +Decision: the real FlashInfer SM70 path has a promising concurrent QSA +component result, but do **not** enable it globally. B1 regresses, QSA indexer +and projection costs are not included, and no E2E/quality gate is complete. +Next: local B1/merge optimization, graph-safe serving workspace/metadata +integration behind opt-in capability detection, then matched quality/latency +validation. Global server settings must not become shortcut routing gates. +The benchmark process exited and GPU 0 returned to desktop-only memory; +no owned worker/API remains resident. + +## Subsequent source-port targets + +The same pinned source already exposes useful next candidates, but backend +names alone do not determine implementation or speed: + +| Area | Source evidence | Required work before a V100 claim | +| --- | --- | --- | +| GDN state update | `flashinfer/gdn_kernels/gdn_decode_pretranspose.py` uses CuTe DSL and cpasync, FP32 recurrent state | Port load pipeline/state layout to SM70, preserve gate math, state-pool read/write indices and graph lifetime | +| Fused convolution/GDN | `flashinfer/gdn_kernels/experimental/kernel/gdn_fused_decode_sm120.cu` is SM120-specific | Extract reusable fusion/dataflow; native SM70 instruction and reduction implementation, not just a capability bypass | +| TP communication | `include/flashinfer/comm/vllm_custom_all_reduce.cuh` explicitly derives from vLLM/SGLang; TRT-LLM fusion variants are separate | Diff actual algorithms, NVLink/P2P graph buffer ownership and barriers; measure message-size crossover and useful fusion, not duplicated old allreduce | +| Projections/HC | FlashInfer has GEMM/communication implementations, not a generic faster replacement for every pointwise Triton kernel | Select by measured operator shape, precision and collective placement; preserve model arithmetic and current native MoE path | + +This work is AI-assisted. Keep implementation and unsuccessful paths in the +owned Draft PR; production admission requires independent correctness and +performance evidence, not compilation alone. diff --git a/docs/design/sm70_qsa_page4_allocation_invariance.md b/docs/design/sm70_qsa_page4_allocation_invariance.md new file mode 100644 index 0000000000..a6e9b35008 --- /dev/null +++ b/docs/design/sm70_qsa_page4_allocation_invariance.md @@ -0,0 +1,201 @@ +# SM70 QSA page4 allocation invariance + +## Defect and contract + +The selected logical tokens and their K/V values can remain identical while +the physical KV page allocation changes between requests. The grouped page4 +planner formerly emitted pages in physical hash-slot order, within active-row +categories. Changing allocation therefore changed the online-softmax and +tensor-core reduction order. This can perturb FP16 attention outputs and +downstream logits without an incorrect selected set or incorrect K/V values. + +The fix guarantees a stable logical plan for fixed query grouping, logical +selections, visibility, and page-alias relationships. Physical addresses are +payload, not ordering keys. It does **not** promise bitwise equality across +different batch shapes, changed prefix-sharing relationships, different +attention implementations, or quantization formats. + +## Implementation history + +| Change | Role in this path | +|---|---| +| [#378](https://github.com/1CatAI/1Cat-vLLM/pull/378), integrated with [#382](https://github.com/1CatAI/1Cat-vLLM/pull/382) | Introduced the SM70 Flash-V100 virtual-page4 path; the single-row path sorted physical microblock IDs for locality. | +| [#387](https://github.com/1CatAI/1Cat-vLLM/pull/387), commit `94ce990ce85abeb12e3948ee1c4f518594bacf25` | Added eight-query K/V sharing, physical-page hash deduplication, mask merging, category packing, and hash-slot-order emission. | +| [#466](https://github.com/1CatAI/1Cat-vLLM/pull/466), commit `186c9e3585b109b88c070603a181dd6826400153` | Lowered the default page4 admission from 4096 to 64 actual query rows and admitted grouped prefixes with XQA tails. It widened exposure, rather than introducing the hash planner. | + +These are source-history findings, not a GPU regression bisection of every +historical release. This defect is independent of AWQ grouped **decode**: +the observed first difference arose during QSA grouped **prefill**, with the +experimental AWQ grouped-decode gate disabled throughout the diagnosis. + +## Causal investigation + +The diagnostic contract used one frozen AWQ checkpoint/runtime, four V100s, +TP4, MTP0, FP16 activations and KV, MRv2, FULL_AND_PIECEWISE graphs, +8192 batched tokens, prefix caching off, asynchronous scheduling off, +fixed prompt token IDs, fixed enqueue order, and fixed request-slot order. +The four prompts contained 42/41/58/44 tokens. Each diagnostic request generated +only its first token, with temperature 0, `min_tokens=0`, and `ignore_eos=false`. +This was numerical diagnosis, not an output-quality or throughput benchmark. + +1. In repeated C4 batches, all four ranks had identical inputs through the first + three layers. The first different boundary was Layer3 QSA (zero-based layer + numbering). Attention output max absolute difference was `0.0009765625`; + full-logit max absolute difference was `0.01171875`, with unchanged argmax. +2. Finer observations found identical Q/K/V, logical positions, request mapping, + selected token indices, and all effective K/V read in logical order. + Physical block tables differed. All 23 query groups retained the same + logical page/mask sets, but 17 had a different order. QSA core max absolute + difference was `0.00048828125` on each rank. +3. Holding physical allocation fixed, 24 extra attention replays were bitwise + identical. The measured symptom was not fixed-input kernel randomness. +4. Diagnostic CPU sorting of only Layer3's plan restored equality there and + moved the first difference to Layer7, the next QSA layer. +5. Applying the same control to all 12 QSA layers made all four ranks' 175 + observed boundaries and complete logits bitwise identical. +6. Removing sorting while retaining observation and CPU synchronization brought + the Layer3 difference back. Extra synchronization alone did not explain the + result. + +The actual C4 step had 185 query rows: 184 grouped rows plus one XQA tail. +The short C1 control had 44 rows and did not enter page4. Long C1 prefill can +still enter the path; concurrency labels are not route evidence. + +Instrumentation can affect compilation boundaries. Causal attribution rests on +same-instrumentation intervention/reversal and local fixed-input replays, not +on interpreting diagnostic timings as performance. The earlier freely batched +70-versus-61-token answer divergence and the independent AWQ W13 operator +microdifference are not claimed to be completely explained by this experiment. + +## Upstream comparison and duplicate-work check + +Original vLLM QSA iterates logical selection positions, then uses the block +table for addressing. It does not use this SM70 physical-hash union planner. +Related upstream work must not be conflated with this defect: + +- [1Cat #394](https://github.com/1CatAI/1Cat-vLLM/pull/394) already provides + exact lexicographic QSA top-k selection on the tested SM70 contract. + Those selections were bitwise equal during this investigation. +- [vLLM #55122](https://github.com/vllm-project/vllm/pull/55122) addresses + generic `persistent_topk` membership/tie/order nondeterminism, not ordering + subsequently introduced by the page4 planner. +- [vLLM #54873](https://github.com/vllm-project/vllm/pull/54873) skips unused + sparse-attention selection entries and tunes launch profiles. It does not + fix the 1Cat planner. +- [vLLM RFC #55394](https://github.com/vllm-project/vllm/issues/55394) proposes + a related query-tile union. Its prototype sorts logical blocks before + physical mapping. That principle is useful here; the GB10 single-request + prototype is not a ready-made SM70 concurrent replacement and is not ported. + +At the 2026-09-04 duplicate check, PR #55122 was open, PR #54873 was merged, and +RFC #55394 was open. No direct repair of this planner was found. The existing +Triton fallback remains a correctness/performance control; upstream use does +not by itself establish V100 performance or cross-batch invariance. + +## Narrow repair + +The grouped planner retains physical-page deduplication and OR-merged token +masks. Each entry also records its smallest logical owner: +`(first contributing query within the group, logical four-token block)`. +An atomic minimum makes shared-page ownership independent of insertion order. +CUB block radix sort orders entries by active-row category and logical owner. +The original category packing, eight-page padding, and attention kernel remain. +This uses CUB's existing sorting primitive, not a new sorting algorithm. + +The single-row XQA path, including non-grouped tails, also needs repair. Its +existing GPU `torch.sort` now sorts packed logical keys carrying physical IDs +as payload. The causal partial page remains after complete pages and invalid +slots remain last. Only integer planning changes; attention arithmetic, +weights, quantization, scheduler policy, and route thresholds are unchanged. + +The grouped hash and owner arrays occupy 96 KiB of shared memory. Once entries +are held in registers, that storage is reused for CUB sorting and category +scans. There is no added global-memory grouped workspace or weight/KV copy. +The single-row sorting keys grow from int32 to int64, so temporary metadata +memory is not claimed to be unchanged. Resource use and latency require GPU +measurement; absence of a global grouped allocation does not imply zero cost. +The tested CUDA 12.8 SM70 binary reports 128 registers per planner thread, +zero local memory, and zero stack bytes. Dynamic shared memory is 96 KiB per +CTA; the resource dump's zero static-shared value does not include it. + +## Validation + +The regression is `tests/kernels/test_sm70_qsa_page4_plan.py`. It checks exact +logical-reference plans, shared physical pages, collisions, invalid rows, +all-empty groups, wide unions, selection permutations, page sizes 4/16/32, +graph replay with relocated inputs, contiguous/interleaved FP16 and E4M3 KV, and a +185-row grouped-plus-XQA-tail batch. + +```bash +.venv/bin/python -m pytest -q tests/kernels/test_sm70_qsa_page4_plan.py +.venv/bin/python -m pytest -q tests/models/qwen4_exp/test_qsa_ops.py +``` + +Before this repair, the initial 20-case regression produced 17 failures and +three passes (the all-empty controls). Both attention relocation checks failed +at the bitwise-output assertion. The later single-row and 185-row integration +tests were added separately and must not be counted as part of that initial run. + +The candidate builds with CUDA 12.8 / SM70. On one V100, all 26 new regression +cases and all 14 existing QSA-ops tests pass (40 total). All applicable +pre-commit hooks pass. Captured Layer3 inputs from all four TP ranks reproduce +the baseline relocation difference and become bitwise equal after the repair: + +| Rank | Baseline different output elements | Baseline max absolute difference | Repaired different elements | +|---|---:|---:|---:| +| 0 | 1625 | 0.00048828125 | 0 | +| 1 | 1852 | 0.00048828125 | 0 | +| 2 | 1515 | 0.00048828125 | 0 | +| 3 | 1866 | 0.00048828125 | 0 | + +The existing Triton fallback is also bitwise allocation-invariant for these +four captured pairs. This is direct replay evidence, in addition to the +upstream source inspection above. + +### Bounded operator timings + +CUDA-event medians, three warmups and 20 samples per call, one V100, FP16 KV, +Hq/Hkv/D = 6/1/256, page size 16. The synthetic cases have four 64K requests, +512 shared selected logical blocks per query, randomized physical allocation, +and the stated total query-row count. They favor K/V sharing and are **not** +a server concurrency contract or full-model throughput measurement. + +| Input | Old page4 (ms) | Repaired page4 (ms) | Change | Existing Triton (ms) | +|---|---:|---:|---:|---:| +| Captured 185-row step, allocation A | 0.3518 | 0.3615 | +2.8% | 2.6552 | +| Captured 185-row step, allocation B | 0.3600 | 0.3635 | +1.0% | 2.4049 | +| Synthetic 1024 rows | 1.5130 | 1.5811 | +4.5% | 14.3713 | +| Synthetic 8192 rows | 9.1136 | 9.7823 | +7.3% | 99.3587 | + +The Triton column comes from the candidate probe. Baseline probe Triton +medians were 2.6563/2.6588/14.4645/99.3930 ms respectively; short host-driven +operator calls show noise and these single-session numbers are not confidence +intervals. For 8192 rows, planner-only time grows from 0.6610 to 1.3251 ms. +The cost is real; retaining page4 with logical ordering is still preferable +to the measured full Triton fallback on this V100 contract. Synthetic +reference errors versus Triton remain small (repaired maximum absolute +difference 0.0001220703125 in both cases), not bitwise cross-kernel equality. + +### Natural-EOS full-model sanity + +A bounded AWQ / TP4 / MTP0 / FP16-KV run used the same frozen compatible +runtime as the diagnosis, including its existing AWQ wrapper admission repair +and disabled experimental AWQ grouped-decode gate. No model trace, sampler +replacement, or attention intervention was installed. Prompt IDs, enqueue +order, and free request-slot order were fixed; each C1/C4/C8 shape ran twice +with temperature 0, `min_tokens=0`, `ignore_eos=false`, and a 96-token limit. + +All 26 request outputs stopped naturally, had finite reported logprobs, and +passed basic answer checks. Token IDs matched exactly between the two runs +of each shape. Worker provenance confirmed the repaired binary and Python +source on all four ranks. Actual page4 route logs reported 184 grouped + 1 +XQA row for C4 and 368 grouped + 2 XQA rows for C8. KV capacity was 386,392 +tokens, unchanged from the prior same-configuration run. + +Cross-shape/position differences **remain**: the same open-ended Chinese +prompt produced 61 tokens at C1, 70 at C4, and 61/71 at its two C8 positions. +Each position reproduced its own sequence on the repeat. This run validates +short same-shape repeatability, not cross-batch or cross-position invariance; +the remaining divergence has not been localized by this model sanity check. +It also does not establish long-context quality, NVFP4 model acceptance, E2E +performance, or a resolution of every prior generation divergence. diff --git a/docs/design/sm70_qwen38_nomtp_concurrency.md b/docs/design/sm70_qwen38_nomtp_concurrency.md index 0564cf14e2..0553c17fb4 100644 --- a/docs/design/sm70_qwen38_nomtp_concurrency.md +++ b/docs/design/sm70_qwen38_nomtp_concurrency.md @@ -1076,6 +1076,1200 @@ All completed GPU workers exit. Current endpoint evidence remains the grouped-MoE candidate's 27.371-ms C16 result, pending model-quality admission; the overall concurrency and C1/quality goals remain **unmet**. +### Fused HC pointwise + disjoint push gather screen (2026-09-05) + +The rejected plain sharding path is **not** promoted. The next prototype +implements the previously identified missing fusion, without production +dispatch or CMake changes: + +- `benchmarks/kernels/sm70_hc_push_gather.cuh` +- `benchmarks/kernels/sm70_hc_push_sidecar.cu` +- `benchmark_sm70_hc_batch_tp4.py --mode fused` + +The down producer applies HC SiLU to each rank's 80 local low-rank values, +publishes only 88 values/row (including injection padding), and gathers the +four disjoint pieces into the 320-wide lora plus rank3 injection output. +The up producer fuses four-branch sigmoid/FMA/mean with publication of each +rank's 640 hidden coordinates, then gathers the final 2560-wide result. +There is no zero-filled full tensor or redundant summation of disjoint data. +Local FP16 GEMMs and their materialization boundaries remain unchanged. +The complete candidate has five kernels instead of the plain sharding's +eight: down GEMM/reduction, down SiLU/gather, up GEMM, mix/gather. + +Both native kernels reuse existing SM70 two-epoch push storage and volatile +16-byte publication/poll/clear helpers. Packing-to-CTA mapping stays at one +pack per thread and 128 threads per CTA, preserving epoch indexing when +interleaved with the existing collectives. No new buffer or communicator +layout is introduced. All lifecycle/method bindings and these benchmark +functions must be compiled in the **same DSO**, never passed a communicator +created by another loaded sidecar. A complete benchmark-only sidecar source +is included for reproducibility; it is not an inference-backend replacement. + +PTX from the retained Triton HC kernels was inspected before implementing +sigmoid: explicit `mul`, `ex2.approx`, `add` and `div.full` preserve the +existing operation sequence, followed by the same ordered FP32 FMA and FP16 +rounding. Signed zero is canonicalized only at the boundaries where the +unfused disjoint all-reduce adds positive zero. NaN sentinels are escaped +using the existing helper. This is not activation or weight quantization. +CUDA 12.8 ptxas reports 40/44 registers for down/up producers, zero spills +and 64-byte stack frames; these are compiler properties, not occupancy or +HBM utilization measurements. + +V100 TP4, actual checkpoint layer-0 attention-HC weights, **16 distinct +weight allocations**, same 16-call graphs and five alternating samples; +all completed timing groups pass foreign-GPU-process checks: + +| M | Replicated HC baseline | Fused sharding candidate | Per-call reduction | +| ---: | ---: | ---: | ---: | +| 4 | 29.320 us | 23.411 us | 20.15% | +| 8 | 30.029 us | 24.270 us | 19.18% | +| 16 | 33.222 us | 26.854 us | 19.17% | + +Every candidate timing sample is below its paired baseline at each width; +raw clock/timing variation is retained, not discarded. Scaling these local +deltas to 96 HC calls gives **0.55--0.61 ms as a screening estimate only**: +these are copies of layer-0 weights and synthetic activations, not an actual +96-HC execution, endpoint saving or a revised 238/420/728-tok/s result. +The prior accepted C16 grouped-MoE candidate remains 27.371 ms pending +model-quality admission. Its gap to 21.978 ms is not closed by this screen. + +At all four ranks and activation scales 0.25/1/3, the fused candidate equals +the **unfused sharded** reference with zero tolerance, and poisoned-output +changed-input graph replay equals candidate eager. Against the original +replicated GEMM baseline, the existing association differences remain: +block/injection relative-L2 up to `2.981e-4`/`4.385e-4`, max-abs +`0.00390625`/`0.0078125`. No full-model quality pass is claimed. This fusion +adds no observed difference beyond the already recorded GEMM sharding, but +that is not sufficient for default enablement. + +The queued `--check-only --tokens 1,2,...,17,24` run extends graph/shape +coverage without collecting irrelevant timing; its result must be checked +before production-wrapper integration. Preparation must happen after real +weight loading and before capture, not in a constructor that only allocates +weights or inside fake-tensor tracing. Preserve the existing M1 path and +prefill/unsupported fallbacks. New native calls must resolve to the same +communicator owner rather than falling back across DSO boundaries. + +Artifact: `.artifacts/raw-audit/hc-fused-push-v1-copies16.{json,log}`. +Test DSO SHA256: +`62d77f78b48837565cd1f132f55766d81ad9747ff65121383eb9782862a9e69d`. +It was built from the task sidecar with the new header; the included portable +sidecar has the same bindings and includes the same native source/header. +For a portable rebuild, in a task uv environment with matching Torch/CUDA, +set `CUDA_HOME`, `TORCH_CUDA_ARCH_LIST=7.0` and a private +`TORCH_EXTENSIONS_DIR`, then call `torch.utils.cpp_extension.load` with +`sources=["benchmarks/kernels/sm70_hc_push_sidecar.cu"]`, +`extra_cuda_cflags=["-O3", "-lineinfo"]`, `extra_ldflags=["-lcuda"]`, +`extra_include_paths=["csrc"]`, and `is_python_module=False`. +Set `VLLM_SM70_CUSTOM_AR_LIBRARY` to that returned library before launching +torchrun; never load a second `_C_custom_ar_flashnext` owner in that process. + +The old undercoverage reproducer also completed. With the old binary and +`QWEN38_BATCH_BLOCKS=1`, all four ranks report **9216/10240 elements (90%)** +unwritten at the 20-KiB case, starting at element 1024. The prior fixed-v4 +64-cycle test with the same override passes. This closes the reproducer/fix +loop; it is not evidence of a bug in the unset/default launch geometry. +Artifact: `.artifacts/raw-audit/smallmsg-undercoverage-old-retry.log`. + +### Private HC channel and independent-QSA follow-up (2026-09-05) + +The latest remote-state audit found #474 **merged** at `cc22156c2b`, not +still Draft. The subsequent `de1754f744` fusion was not in that merged head. +Continue in Draft #504, branch +`codex/v100-flashnext-batch-hc-qsa-followup-20260905-082250`, worktree +`worktrees/v100-flashnext-batch-hc-qsa-followup-20260905-082250`, based on +`755baae1d075ee04fa9096b23fc0225b23589a86` (includes merged #481). +The fusion-only commit was cherry-picked as `0936cd23ca`; no stale PR edit, +force-push or automatic performance rebaseline was performed. All following +retained HC numbers use the **old frozen source and DSO**, not this new main +or a clean wheel. Full-model quality and 238/420/728 targets remain unmet. + +The formerly queued width check completed: M1 through M17 and M24, all four +ranks and three activation scales, poisoned graph replay equal to eager and +zero-tolerance equality against the unfused sharded implementation. This +does not erase the original replicated-GEMM association differences. +Artifact in the previous worktree: +`.artifacts/raw-audit/hc-fused-push-v1-width-check.{json,log}`. + +The fused benchmark now allocates one **private HC communicator** per TP +rank, separate from the auxiliary-stream ordinary/sum2 communicator. Both +are owned by the same DSO. This is a conservative prototype isolation gate, +not evidence that a shared-channel race occurred. Do not allocate one such +communicator per HC layer in production: metadata/workspaces (including an +8-MiB rank-data tensor) cost real memory. Reuse existing dedicated channel +concepts from #481 when designing production admission, with native owner +and buffer ABI checks; do not mix lifecycle calls between different DSOs. + +Auxiliary stress: M4/M8/M16, 64 cycles each, changed HC and sum2 inputs, +rank-skewed enqueue order on two streams, 16 sum2 operations per graph. +HC remains zero-tolerance equal to unfused sharding and auxiliary sum2 is +exact at all four ranks. This is operator/state evidence only. Artifact: +`.artifacts/raw-audit/hc-fused-private-aux-v1.{json,log}`. That first auxiliary +run did not restore the timing input after stress, so its timing is not used +as the standard performance contract. The benchmark now restores the input. + +Separate no-aux standard-input measurement, 16 rotating real-weight copies, +16 HC calls per graph, five alternating groups, per-group maximum TP rank: + +| M | Replicated HC | Private-channel fused HC | Reduction | +| ---: | ---: | ---: | ---: | +| 4 | 31.373 us | 26.803 us | 14.57% | +| 8 | 32.150 us | 27.346 us | 14.95% | +| 16 | 34.197 us | 28.754 us | 15.92% | + +Artifact: `.artifacts/raw-audit/hc-fused-private-perf-v2.{json,log}`. +DSO SHA256 remains +`62d77f78b48837565cd1f132f55766d81ad9747ff65121383eb9782862a9e69d`. +These are the conservative current HC micro figures. The 96-call projection +is only 0.44--0.52 ms, not an endpoint saving. Different processes/clocks do +not permit attributing the entire difference from the earlier 19--20% +screen to channel isolation. No new production hook/default is enabled. + +**QSA import/admission audit:** the retained source-build launcher appends +an obsolete `v100-qwen38-exact-decode80-20260829-030630/flash-attention-v100` +source tree that has no native `.so`. CPU-only reproduction matches the +import failure recorded at line 424 of `grouped-runtime/control_before.log`. +A package-scoped bootstrap successfully imports the old task's built +Flash-V100 package (grouped ABI 2), without broadly inserting `build/lib` +and changing other dependencies. Native SHA256: +`daa665b9e81914de1b477e278524fccdfd45d39f09193f67268f48604ac92b49`. + +Crucially, `_use_sm70_qsa_xqa_page4` defaults to **rows >=64** for FP16 KV, +and the frozen launcher does not override it. The import warning is from +prefill, not proof that C4/C8/C16 decode lost an admitted native route. +Do not present this as the root cause of slow batch decode or silently +lower the threshold. `benchmark_sm70_qsa_batch_routes.py` instead screens +forced Triton, direct XQA and padded grouped Page4 for independent requests, +including padding/planning/copies, random physical pages, six changed-input +and slot-map graph replays, and an explicit FP32 selected-attention oracle. +Grouped reuse must not assume adjacent rows share request KV. The existing +MTP benchmark's `xqa` wrapper can itself select grouped mode for M>=8; the +new screen calls each native route directly to keep labels unambiguous. + +The first native micro reached the GPU, then failed M1 FP32-oracle tolerance +(max reported failing difference `0.00308471`, versus `atol=0.002`). No timing +was admitted. Source audit then found its reused MTP benchmark generator +always appends the last three tokens. That is **not the production QSA +contract**: the compressed selection has 512 complete Page4 blocks plus +only `visible_length % 4` open-tail tokens and `-1` padding. At length 8192 +there is no open tail, and direct XQA correctly excludes that noncanonical +extra data. This failed input is not evidence of model-quality degradation. +The first log did not label the failing route, so do not assert it proved +a specific native kernel error. Artifact: `.artifacts/qsa-native-frozen-v1.log`. + +The old generator is now corrected for both independent decode requests and +causal MTP rows (shared pages must be complete for the earliest query). +New CPU metadata coverage: **24 passed** across M4/M5/M8/M16, independent +versus single-request rows, and three overlap levels. Existing current-main +QSA launch/route tests: **27 passed, 1 GPU-only skip**. These are metadata +checks, not GPU/model quality. The new native screen uses the production +expansion operator, independently checks its output against Torch integer +arithmetic, masks invalid indices in its FP32 oracle, and sweeps all four +tail residues. It restores the declared 8192 length for timing. Tolerances +are unchanged; failed routes retain per-route diagnostics and no admitted +timing. The redundant v2 diagnostic waiter was stopped before GPU launch; +v3 canonical run is queued behind verified foreign GPU ownership. +Do not edit its live shell launcher or restart it merely because it waits. + +New-base sidecar compile/import passed (CUDA 12.8 / Torch 2.10, SM70, +`-O3 -DNDEBUG -lineinfo`), SHA256 +`872f4e3a37cdf57c2e5ddd010c7b95cd968e4f8d6f5a64fd79fff87063c51f8a`. +It resolves lifecycle plus batch HC methods to the same sidecar namespace +and reports the new 720000-byte push allocation. **Not yet GPU tested**. +This benchmark-only sidecar does not expose #481's M1 HC methods; do not +use it for an endpoint/C1 comparison or production deployment. The future +integration must preserve all existing native features and separately +validate the loaded owner's capabilities. All prior private-HC GPU evidence +remains explicitly tied to the frozen old DSO rather than this new binary. + +New-worktree CPU test logs: `.artifacts/qsa-cpu-tests.log` and +`.artifacts/qsa-benchmark-indices-tests.log`. Canonical GPU result locations: +`.artifacts/qsa-native-frozen-v3-canonical.{json,log}` (inspect completion, +not file presence, before citing results). Latest endpoint and model-score +status are unchanged; do not convert this metadata repair to throughput. + +### Canonical QSA outcome and HC weight-view screen (2026-09-05) + +The canonical QSA job completed (exit 1 because grouped replay admission +failed), with exclusive single-GPU timing. No full-model service was launched. +FP16 KV, independent 8K requests, GQA6/D256, 512 selected compressed blocks, +five alternating timing groups, 16 calls per graph, 40 replays per group: + +| M | Frozen Triton | Direct XQA | Padded grouped Page4 | +| ---: | ---: | ---: | ---: | +| 1 | 26.610 us | 221.773 us | 381.122 us | +| 4 | 53.162 us | 283.080 us | 1738.734 us | +| 8 | 98.955 us | 287.075 us | Not admitted | +| 16 | 166.763 us | 289.534 us | Not admitted | + +All three routes pass the unchanged FP32-oracle tolerance at every shape and +all six changed-input/slot-map cycles. Direct XQA and Triton graph outputs +equal their eager outputs. Grouped M8/M16 graph outputs are not bitwise +equal to eager, although their FP32 relative-L2 remains below `3e-4`. +Do not call this model-quality corruption or infer its cause from the small +operator differences. Grouped's ordering/reduction behavior is not localized; +it has no admitted timing at these widths. The slow M1/M4 result is already +sufficient to reject this native small-batch route as the next optimization. +No change to the rows>=64 production admission threshold follows. The source +is the frozen old QSA implementation and the pinned Flash-V100 binary, not a +claim about every Flash-V100 build or its existing large-prefill path. +Artifact: `.artifacts/qsa-native-frozen-v3-canonical.{json,log}`. + +HC fusion now also passed a GPU screen against the **new main-base sidecar** +(`872f4e3a37...`), rather than merely importing/compiling it. Same real layer-0 +weights, 16 distinct weight allocations, five alternating paired measurements, +16 complete HC calls/graph, all four TP ranks. The benchmark includes 64 +changed-input auxiliary-stream sum2 cycles per width, then restores the +original timing input. All fused-versus-unfused sharded and graph-versus-eager +checks pass with zero tolerance. The original replicated-GEMM association +differences remain; there is still no model-score admission. + +| M | Packed: baseline / fused | All views: baseline / fused | Down view: baseline / fused | +| ---: | ---: | ---: | ---: | +| 4 | 31.406 / 26.898 us | 31.389 / 30.707 us | 31.294 / 26.989 us | +| 8 | 32.214 / 26.626 us | 31.274 / 29.830 us | 32.059 / 27.309 us | +| 16 | 33.819 / 27.885 us | 31.723 / 29.723 us | 32.410 / 27.000 us | + +The three layouts are separate processes, each with its own paired baseline; +do not attribute every absolute-time difference to layout. Packed retains +14--18% local savings. Fully strided `torch.bmm` retains only 2--6% local +savings, so it is not the performance choice despite eliminating the extra +weight storage. Down-view retains about 14--17% versus its original HC +baseline, but these data do not prove <=1% regression against the packed +candidate at every shape (in particular M8). Keep that distinction rather +than claiming a free speed improvement. Its 96-call projection is only +0.41--0.52 ms, not a new endpoint result. + +The view design uses ordinary checkpoint FP16 values: down is a contiguous +88-row slice (only rank3 uses the four injection values in its extra rows), +up can be a four-branch strided batched GEMM with shared lora input and +disjoint output columns. This applies the tensor-contraction layouts described +in [NVIDIA's strided batched GEMM reference](https://developer.nvidia.com/blog/cublas-strided-batched-matrix-multiply/); +it introduces no weight or activation quantization. The new CPU layout suite +passes **16 cases** (four ranks x four widths), checks no output-cell overlap, +storage aliasing, and visibility of original-weight updates. It is not a GPU +arithmetic oracle or a model-quality score. + +Measured additional shard storage for 16 HC weight copies is 52.5 MiB +(packed), zero (all views), and 25 MiB (down view). For 96 HC pairs these +amount to 315/0/150 MiB respectively: down-view avoids **165 MiB/worker** of +the prototype's additional weight copies. Original full weights, private +communicator and graph workspaces remain allocated; this is not total model +memory usage. The benchmark reports aliased-view bytes separately from +additional allocated shard storage. Artifacts: + +- `.artifacts/hc-main-{packed,views,down-view}-v1.json` +- `.artifacts/hc-main-layouts-v1.log` +- `.artifacts/hc-main-down-view-v1.log` +- `.artifacts/hc-weight-views-cpu.log` + +All completed task GPU workers exited; the API was not started. The next +production work must preserve #481's M1 native capabilities, use an isolated +HC channel owned by the same DSO, prepare any remaining up-weight copy only +after real loading and before capture, and retain dynamic prefill/unsupported +fallbacks. Do not globally bind max-seqs/chunk/KV or silently replace the +frozen 81-tok/s C1 contract. After route-hit/quality checks, measure the actual +engine and datasets. Current endpoint evidence remains C16 584.568 tok/s / +27.371 ms, with overall throughput and model-quality gates **unmet**. + +### Opt-in HC runtime integration and component gate (2026-09-05) + +The successful down-view/packed-up variant now has an actual model dispatcher, +behind `VLLM_SM70_QWEN38_BATCH_HC_FP16=0` (default off). The tested CUDA kernels +are shared by production custom-AR bindings and the benchmark. A complete +native-owner sidecar also preserves existing M1 bindings; the earlier +benchmark-only sidecar that omitted them must not be used for model runs. + +One isolated communication channel is created per eligible TP group, not per +layer or hot-loop invocation. The up-weight shards are nonpersistent buffers +prepared after quantization post-load hooks, with pointer-preserving reload. +The down shard aliases original FP16 storage. FP16 dimensions and native +TP4 connectivity are local capabilities; no KV, scheduler chunk, max-seqs or +checkpoint-ID gate is added. Unknown quantization/LoRA/offload layouts fall +back. The opaque op makes the decode-width decision at runtime and preserves +the previous M1 fused-HC delegate and prefill behavior. + +Validation from this worktree, CUDA 12.8 / Torch 2.10.0+cu128: + +- `tests/models/qwen4_exp/test_sm70_batch_hc.py` plus + `tests/quantization/test_sm70_nvfp4_grouped_decode_dispatch.py`: **60 passed**, + `.artifacts/hc-production-cpu-v2.log`. +- Staged pre-commit hooks all passed after correcting formatting and using + accelerator device/synchronization APIs. Artifact: + `.artifacts/hc-production-precommit-v3.log`. +- Real `GatedResidual` and vLLM TP4 communicator on GPU 0--3: + M4/M8/M16 hit the candidate; M1, short prefill, M17 and prefill M32 correctly + fall back. Fallback output equals the original route. Eight changed-input, + poisoned-output graph replays match eager. Prefill-first dynamic compile, + re-preparation pointer stability and communicator destruction passed. +- Actual layer-0 weights, random activations: candidate/reference block max + absolute difference 0.00048828125 and injection 0.001953125 for M4/M8/M16. + These are not dataset-quality guarantees. No additional quantization used. + +Runtime artifacts: `.artifacts/hc-production-runtime-v1.{json,log}` and +`.artifacts/run_hc_production_runtime.sh`. Complete sidecar: +`.artifacts/torch-extensions/vllm_sm70_hc_batch_production_v2/` +`vllm_sm70_hc_batch_production_v2.so`, SHA256 +`3664deec1c713e3c4e0fe2bb5de22cc783a0eb99ebb7539a0376c39c270bdab2`. +This is source plus a native sidecar, not a clean release wheel. + +Next full-model comparison retains the old fixed-width runner unchanged, +with both arms on current task source, complete sidecar and the same pinned +valid Flash-V100 package. Grouped MoE is fixed on; HC is the only arm switch. +The corrected Flash-V100 import means this is a new controlled comparison, +not a relabeling of the old exact baseline. First run the first 16 GSM8K test +items with official sampling and natural EOS (16K maximum), then the original +deterministic 8K/256 C1/4/8/16 speed probe. This small health screen is not +final dataset/tool/schema/PPL admission. End-to-end and quality targets remain +unmet pending results; no production default changed. + +### Concurrent API quality coverage and upstream review (2026-09-05) + +The existing BFCL/JSONSchemaBench API runner sends requests serially. It cannot +prove that a new batch-only HC path preserves tool/schema quality. Add +`benchmarks/benchmark_sm70_batch_tool_quality.py`, reusing the existing BFCL +name/argument scorer, JSON Schema normalization/validation and SSE collector. +The fixed subset is 16 entries from each of four BFCL categories plus 16 +size-stratified WashingtonPost schemas (80 total), unchanged between arms. +Sampling defaults to temperature 1, top-k 20, top-p 0.95, natural EOS and 16K +maximum. Thinking is explicitly selectable and must match between arms. +No tools are executed; this is a scored first-turn API subset, not multi-turn +agent task completion or the official leaderboard. + +The client retains input manifests/source hashes, per-case seed, complete +payload/SSE/outputs, transport failures without retries, and request overlap. +Incomplete streams and length-truncated JSON cannot pass just because a JSON +fragment parses. Observed client concurrency is not evidence of actual GPU +batch width: worker dispatch/trace remains required. Run as a module: + +```bash +python -m benchmarks.benchmark_sm70_batch_tool_quality \ + --base-url http://127.0.0.1:18184 --model FlashNext \ + --bfcl-dir "$BFCL_DATA" --schema-dir "$JSONSCHEMA_WASHINGTONPOST" \ + --concurrency 16 --output "$RESULT_JSON" +``` + +Use the task virtualenv's Python, not system Python. The `--dry-run` CPU pass +selects all 80 real cached cases and validates their schemas. Client unit +tests verify actual overlap at C1/4/8/16, stable case/seed order, retained +transport failures, malformed tool names and truncated/schema-invalid output: +12 passed. These are harness checks; API quality scores remain pending. + +The full-model speed/health task waited on confirmed foreign workers and +reservation ownership rather than preempting them. While still queued, its +first waiter was terminated deliberately to fix a CPU-reproduced tokenizer +contract: `apply_chat_template` returns `BatchEncoding` by default here. +Explicit `return_dict=False` and a one-time vocabulary-size lookup now validate +all 16 GSM8K prompts before model construction. The repaired CPU preflight +passes; no model startup was spent on either harness repair. A single +replacement control waiter is retained, not duplicate GPU jobs. + +Primary-source review: + +- [vLLM/HPC-Ops low-latency MoE](https://vllm.ai/blog/2026-07-06-vllm-hpc-ops) + motivates indexed input reads, shared routing/task maps and occupancy-first + scheduling. Our grouped native-NVFP4 path already reads original tokens by + route and shares groups between W13/W2. Its static grids still reserve + worst-case route tiles, a possible next screen after the current full-model + comparison. This is a hypothesis, not a measured speedup. Hopper FP8/PDL + mechanisms cannot be transplanted to SM70; retain TurboMind/native NVFP4. +- [AMoE asynchronous expert serving](https://arxiv.org/abs/2505.08944) + trades queuing/rebatching against latency. It is not a shortcut to the fixed + C1/4/8/16 interactive-step targets, and is not selected for this patch. +- PR #506 was inspected: its remaining norm-prefetch work targets single-token + HC. Keep this comparison frozen; do not mix it into a batch-fusion A/B. + +### Completed HC full-model A/B (2026-09-05) + +Both finite runs completed at source +`b64dae6e5b7d14eb4f0e0861d31600fc866c4e48`. All task model workers exited. +The control/candidate use the same complete HC native-owner sidecar and pinned +Flash-V100 library described above. Native NVFP4 target, FP16 KV/activations, +FP32 GDN state, TP4 V100-SXM2-32GB GPU 0--3, Torch 2.10.0+cu128, CUDA 12.8, +driver 580.173.02, MRV2/no MTP, Prefix Cache + Mamba align, max context 262144, +chunk 2048, max-seqs 16, memory utilization 0.90, full decode graphs. Grouped +MoE stays on in both arms. The only candidate switch is batch HC. + +Original deterministic speed workload: independent 8192-token inputs, forced +256 output tokens, temperature 0/top-p 1/top-k -1, same prompts/seeds and +fixed-live-width engine-timestamp interval selection (head/tail 8 excluded). +The client receive-blocking wait is not used as complete TPOT. + +| C | HC off ms | HC on ms | Off tok/s | On tok/s | On gain | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 | 11.357011 | 11.356076 | 88.051 | 88.059 | 0.008% | +| 4 | 18.753086 | 18.476128 | 213.298 | 216.496 | 1.499% | +| 8 | 22.144332 | 21.970248 | 361.266 | 364.129 | 0.792% | +| 16 | 27.400743 | 27.125071 | 583.926 | 589.860 | 1.016% | + +This is one A/B, not a repeated confidence-interval admission. Real per-step +savings are only 0.277/0.174/0.276 ms at C4/C8/C16, below the component +projection and far short of the 238/420/728 goals. Worker logs confirm all +96 prepared HC up shards and native decode capture at M2/4/8/16, with original +M1 fused up/mix/gather preserved. Candidate loading uses approximately 150 MiB +extra per rank; reported KV pool 541053 -> 529622 tokens (not a 256K quality +test). No default changed. + +The separate natural-EOS first-16 GSM8K health screen has identical prompt IDs +and official sampling (temperature 1/top-p .95/top-k 20, seeds 20260905+i, +thinking enabled, max 16384). Both score **15/16**, both miss item 12 (12 years +versus the reference's 13 years), and neither truncates. Responses have +different lengths/text; this small score tie does not establish non-inferiority +on the required coding/tool/schema/PPL gates. Its duration is not throughput: +61.24 vs 30.42 seconds includes different generated lengths. The legacy +speed JSON `load_seconds` field wraps startup **and this health check** in the +adapter; it is not a model-loading-only or prefill measurement. + +Even the M1 forced-speed completion differs at token 12 despite identical M1 +dispatch; this synthetic continuation runs past natural EOS. Keep this as a +diagnostic, not an automatic model-quality failure or attribution to HC. +PR #494 documents a separate allocation-sensitive QSA planner order issue; +it was inspected, not imported or proved to explain this run. Do not change +the frozen A/B or impose cross-batch bitwise equality instead of score gates. + +Runtime limitation: both arms' frozen base `_C` library lacks the newer +single-token `nvfp4_qwen38_w13_fused_swiglu_out` and +`nvfp4_qwen38_w2_direct_reduce_out`. Source audit places those gates only in +the direct M1 MoE branch; current M4 direct-batch fusion and M8/M16 grouped +routes do not depend on them. Both arms emit the same fallback warnings. +Thus the HC A/B is controlled, but is **not** a clean latest-main wheel or +complete native-optimization baseline. C1's improvement over the old 81-tok/s +run must not be attributed to the new batch HC switch. + +Artifacts in this worktree: + +- `.artifacts/hc-model-runtime/{control,candidate}-v1.json` and matching `.env`; + JSON SHA256 respectively + `c8e13ba1d6de30fa5c362b171e58fcf6370074f5254aad0a706c785538c8bd3f`, + `6cad83b2948a011e526763762afa4ecde7f29eb374fa4bad4e81151fcf7593bc`. +- `.artifacts/hc-model-runtime/{control,candidate}-quality-v1.json`; + SHA256 respectively + `a639cfbe6fbeb9f24f4b577799211bc3b8bdbb1b37643ef9922d392403d9c41e`, + `b7b8e44a685922555bfe4fe27d79f0d7153a720316550b3db6886f57e2f5c7c1`. +- `.artifacts/hc-model-control-v1-requeue.log`, + `.artifacts/hc-model-candidate-v1.log`, + `.artifacts/run_hc_full_model.{sh,py}`. + +Next bounded job: one local candidate API at port 18184, concurrent 64 BFCL + +16 JSONSchemaBench cases, then an explicitly deterministic, prefix-warmed +C16/8K short CUDA graph-node trace. CUDA profiler delay 8/max 16 steps avoids +profiling startup; the profiled request is not accepted throughput. Nsight +2022.4.2 uses `--cuda-graph-trace=node`, not unsupported `node:host-only`. +Task launcher `.artifacts/run_hc_api_quality_trace.sh` waits for verified idle +GPUs/locks, refuses port/artifact overwrite and shuts down its own API by PID +plus process-start generation. It is a finite quality/profile job, not a +resident API. No API quality or new trace result exists at this update. + +### API/trace launcher repair before weight loading (2026-09-05) + +The first API/trace job obtained the GPUs, then failed before worker/model +weight loading. The task-only `hc_api_entry.py` executed `runpy.run_module` +at import time. Multiprocessing `spawn` re-imported it as `__mp_main__`, +started a second API in the child, overwrote the task PID file, and raised +the explicit bootstrapping-phase error. This is a test-launcher failure, not +a model-output result. The job ended with client/nsys status 1; both API PIDs +were confirmed gone and no task-owned GPU worker remained. No quality cases +or trace timing from this attempt are admitted. + +Keep `.artifacts/hc-api-trace/{server-v1.log,client-v1.log,api.pid}` and the +original launcher as evidence. The repaired entry +`.artifacts/hc_api_entry_v2.py` guards both CLI invocation and PID writes +under `if __name__ == "__main__"`. CPU `runpy.run_path(..., +run_name="__mp_main__")` verifies that import leaves argv unchanged and +creates neither API nor PID-file side effects. The guarded V2 launcher uses +the separate `.artifacts/hc-api-trace-v2/` output directory, retains +PID-generation cleanup and port/artifact refusal, and waits on both locks and +real compute processes. Exactly one replacement job is queued after the +first was proven terminal. It polls at 2 seconds with 20-second status output; +no occupied GPU is preempted. No V2 quality/trace result yet at this update. + +The concurrent API scorer also now rejects a response that contains tool +calls but reports a non-tool `finish_reason`, even if its name/arguments are +otherwise correct. The existing 12-test CPU suite covers this added negative +case and passes. This repairs a validation gap, not evidence that the current +API actually emits that defect. Do not substitute a prefill-only PPL run for +decode-path quality: the new HC/grouped-MoE candidates explicitly fall back +on prefill and such a run would not exercise their changed arithmetic. + +### Concurrent API results and measurement repairs (2026-09-05) + +The V2 API job completed with client/nsys status 0. All 80 requests returned +HTTP 200 and ended naturally (53 `tool_calls`, 27 `stop`); client peak inflight +was 16. This proves neither fixed GPU batch width nor model-quality parity. +The same 64 BFCL / 16 JSONSchemaBench cases used temperature 1, top-k 20, +top-p 0.95, per-case seeds starting 20260905, thinking off and max output +16384. Do not compare to a thinking-on control or use request wall duration +as pure-decode throughput. + +| Suite | Correct / total | +| --- | ---: | +| BFCL simple Python | 14 / 16 | +| BFCL parallel | 12 / 16 | +| BFCL multiple | 14 / 16 | +| BFCL irrelevance | 11 / 16 | +| JSONSchemaBench WashingtonPost | 16 / 16 | + +The original BFCL score was 49/64. Offline audit corrected **two scorer false +negatives**, not model outputs: `multiple_8` budget and `multiple_9` gradeDict. +BFCL dictionary ground truth encodes per-key acceptable alternatives (e.g. +`{"min": [300000]}` permits the scalar `300000`). The old helper compared +the scalar to the literal list. Match dictionary and ordered list-of-dicts +values using the rules of the pinned official BFCL `dict_checker` / +`list_dict_checker`, including missing optional keys and literal array +alternatives. Other scoring rules remain unchanged; this is still a subset +scorer, not the full official leaderboard evaluator. + +The pinned reference is BFCL `f7cf7359b7ac615a0b294831c5ba2bc95ee4a000`, +`berkeley-function-call-leaderboard/bfcl_eval/eval_checker/ast_eval/ast_checker.py`. +Eight real-parameter/negative-mutation vectors agree with its actual +standalone `dict_checker`; the expanded CPU client/scorer suite passes +**24 tests**. A plain pytest launch initially imported the source tree's +unbuilt `_C` via root conftest; the CPU-only command uses +`python -m pytest --confcutdir=tests/benchmarks +tests/benchmarks/test_sm70_batch_tool_quality.py -q`. It does not initialize +GPU fixtures or claim GPU validation. + +The corrected BFCL score is **51/64**, with all generated responses, seeds, +payloads and dataset hashes unchanged. The remaining 13 misses are retained +for a matched original-production control, including omitted/wrong arguments, +wrong call counts, name/choice errors and irrelevant tool use. Candidate-only +scores do not establish non-regression. No optimization is newly defaulted. + +Artifacts: `.artifacts/hc-api-trace-v2/candidate-tools-v1.json` (immutable raw), +`candidate-tools-rescored-v2.json`, `.artifacts/rescore_hc_api_tools.py`, +`.artifacts/bfcl-rescore-oracle-v2.log`, `.artifacts/bfcl-dict-scorer-tests.log`. +The rescorer refuses changed dataset hashes, case ordering or request payloads. + +The accompanying trace is **rejected as decode evidence**. Its NVTX labels +all contain prefill (2044--2048 context tokens); the SQLite has 198412 kernels +and zero graph-node kernels. Repeating the 16 prompts did not yield prefix +cache hits, and the fixed 8-step delay captured mixed prefill, not C16 decode. +Retain the qdstrm, manually imported nsys-rep and SQLite in the V2 directory; +do not feed these into a C16 GPU-cost table or repeat that priming strategy. +The API and all four owned model workers have exited. + +Replacement task `.artifacts/run_hc_fixed_decode_trace.{sh,py}` preserves the +frozen independent 8K/256 deterministic speed workload and runtime flags. +It arms profiling only after eight consecutive engine outputs report all +16 running, zero waiting, one emitted token per request, no prefill and no +finished request. Then it captures 16 worker steps after a two-step delay. +CPU tests reject each mixed-width/prefill/finished trigger and verify spawn +import has no launch side effect. Output is isolated in +`.artifacts/hc-fixed-decode-trace/`; the launcher checks actual idle GPUs and +ownership locks, refuses artifact overwrite, and the finite runner explicitly +shuts down its engine. The new trace still requires NVTX/graph-node validation +before attribution; its profiled throughput is never accepted endpoint speed. + +### Valid C16 graph attribution and QSA grid screen (2026-09-05) + +The decode-triggered job completed successfully and its engine, workers and +launcher all exited. NVTX now contains exactly 16 +`execute_context_0(0)_generation_16(16)` ranges per TP rank. All 64 complete +graphs contain 2211 nodes. The source/runtime is the frozen `0ed451fbda` +candidate; subsequent scorer/docs commits do not change `vllm/` or `csrc/`. +This is the same pinned native stack, not a new clean-wheel speed claim. + +Use `.artifacts/attribute_hc_c16_complete_graphs.py` to group by the actual +`(globalPid, correlationId)`, check uniform node counts and drop each rank's +first/last graph. Do not hardcode the old pre-grouped graph's 1976 nodes. +The bundled generic per-token parser is also retained, but its host replay +intervals cut across some GPU graphs and it labels the shared cuBLAS signature +as LM-head. Neither artifact should be presented as an additive wall table. + +For the 56 middle complete graphs, average per-rank graph envelope is +**26.789 ms**, activity union **24.895 ms**, summed kernel service **27.053 ms**, +internal gaps **1.894 ms**, overlap **2.158 ms**. Per-step max-rank envelope +averages 26.799 ms. These are profiled graph diagnostics, not new endpoint +measurements; the accepted unprofiled C16 candidate remains **27.125 ms / +589.860 tok/s**, with the 21.978-ms target still unmet. + +| Source-attributed family | Mean service ms / rank / graph | +| --- | ---: | +| Routed MoE W13 | 4.301 | +| Routed MoE W2 | 2.284 | +| MoE grouping and weighted reduction | 0.268 | +| QSA scoring, top-k and sparse attention (excluding projections) | 3.824 | +| GDN qkvz projection | 1.653 | +| GDN b/a projection plus reduction | 0.583 | +| GDN recurrence / convolution | 1.520 | +| HC down/up projection plus down reduction | 1.947 | +| HC fused pointwise / TP gather | 1.437 | +| HC remaining postops | 0.536 | +| Other TP reductions | 1.195 | + +Other projections, shared expert work and 1.702 ms of unattributed small +kernels remain explicitly in the complete JSON. Role inference uses shapes +and same-stream neighbours, not invented module NVTX labels. The graph +excludes the subsequent LM-head/sampler; do not subtract its envelope from a +different run's endpoint time to manufacture a host-cost closure. + +Within QSA, the scorer costs about 103 us/layer and sparse GQA about +169 us/layer. The scorer grid is `(16, 1026)`, driven by the fixed 256K +capacity (65660 compressed columns) rather than visible length. At this +8K/256 workload, only approximately 512--528 of 16416 CTAs can have live +tiles: over 96% immediately return. This is a source/shape deduction, **not** +an NCU occupancy or measured bandwidth percentage. The native small-batch +Page4 alternative was already rejected; do not retest that dead end. + +New benchmark-only `benchmark_sm70_qsa_scorer_grid.py` tests the existing +scorer at rows 1/4/8/16, lengths 8K/128K/256K and fixed maximum capacity, +with randomized physical pages. The indexer has **four replicated heads**, +not four heads divided by TP4. All 12 cases pass four changed-input, +shrinking/growing/mixed-length poisoned CUDA Graph replays with exact live +scores against the production kernel. This is operator evidence, not model +quality admission. The graph grouping screen uses five alternating-order +timing samples and 20 replays per sample. + +Contiguous grouping is not a safe default: at C16/8K, group 4 improves +65.331 -> 48.179 us, but at C16/256K it regresses 1094.042 -> 1113.856 us +(1.81%). Larger grouping also reduces useful short-context CTA parallelism. +Reject an unconditional `tiles_per_program` change. These hot-cache single- +GPU times must not replace the TP4 full-model trace's absolute scorer cost. + +The next benchmark-only variant, `sm70_qsa_strided_scorer.py`, bounds the +grid while visiting **all** live tiles by a grid-stride loop. It retains each +tile's dot/head reduction and existing validity masks. This follows the +portable scheduling concept in the [Triton persistent-kernel tutorial](https://triton-lang.org/main/getting-started/tutorials/gluon/persistence.html), +not its Hopper-only TMA machinery. PR #507 targets single-row QSA/router; +its body/source were checked and do not cover this batched scorer grid. +Production dispatch is unchanged. The initial queued launcher was cancelled +before acquiring GPUs after a syntax preflight caught a missing closing +parenthesis in the copied benchmark kernel; no GPU result was produced by +that attempt. A syntax-checked replacement waits behind verified foreign +paper-campaign GPU processes and their ownership reservation. + +Artifacts under `.artifacts/hc-fixed-decode-trace/`: raw qdstrm SHA256 +`b2223fcf637e9706f80d959f7184e0a8c2bed5b3a514165bd32f2b8f10d2f7f7`, +SQLite `1815931586378c3fc23d460107e1d7e0dcacdf1aecccbd5ac6483535518ef83f`, +`whole-graph-attribution.json` +`ca5c0b8144a7c587c3d32339435277a7dcc0502c6c14138dfc42833080e2053a`. +Scorer screen: `.artifacts/qsa-scorer-grid/grouping-v1.{json,log}`; +the replacement strided-grid job has its own `strided-v1` outputs and +`.artifacts/qsa-scorer-strided-launch-v2.log`. All benchmark processes are +finite; no API is left resident. + +### Grid-stride scorer follow-up: short gain, long gate still fails + +Both finite strided-grid jobs subsequently acquired idle GPUs, completed all +12 row/length cases, passed the four exact changed-input/length graph replay +checks per case, and released their CUDA processes. The first used the +committed benchmark at `3f1fa2d06e`; the second additionally records compiled +resources and tests a one-stage inner loop while preserving the baseline's +two-stage setting. There is no resident model/API or queued GPU task left +from these screens. + +The two-stage strided-64 C16/8K scorer improves **65.229 -> 44.954 us**. +However its C16/256K result is **1105.306 -> 1174.477 us** (6.26% slower). +Strided-128 reduces that long penalty to 2.26%, still outside the 1% gate. +These are same-process micro pairs, not an end-to-end gain. C1/long and C4/long +regress more; a blanket smaller grid would sacrifice existing useful paths. + +The final one-stage screen does not rescue the long boundary. C16/8K +strided-64 is **65.075 -> 44.493 us**, but C16/256K strided-256 is +**1104.128 -> 1118.720 us** (1.32% slower). Compiled resources for the C16 +baseline/strided kernel are **40960/40960 shared bytes**, **224/238 registers +per thread**, zero spills, two warps. Reducing the requested loop stages did +not reduce the shared-memory footprint. Do not attribute the regression to a +proved shared-memory occupancy change: the compiled limits do not show that. +These are compiler resource facts, not NCU achieved-occupancy/HBM counters. + +Reject promotion/default enablement of both fixed contiguous grouping and +the fixed bounded-grid variant. The ~20-us short scorer saving is only about +0.24 ms across 12 layers even before endpoint validation, far short of the +remaining C16 5.15-ms gap. Next prioritize the **6.85-ms routed MoE family** +and the 2.0-ms sparse-attention forward within QSA, using this new graph +attribution. Do not keep sweeping scorer parameters or reload the full model +to chase this small, currently non-admitted result. Any future scorer work +must address long-grid scheduling while retaining its actual live-length +coverage and existing C1 path. + +Artifacts: `.artifacts/qsa-scorer-grid/strided-v1.{json,log,env,source-sha}` +(JSON SHA256 `fb4600da38fbd4e64ffebd79cfa77bd2557a556d1af61bde7b0d0a6c7dd44bc8`), +`strided-stage1-v1.{json,log,env,source-sha}` and corresponding task launchers. +The stage-1 source is the resource-reporting benchmark change accompanying +this worklog; it does not modify any runtime module. Full C1/C4/C8/C16 +performance admission and original-production quality comparisons remain +open. The owned PR remains Draft and all production defaults are unchanged. + +### MoE reuse follow-up: locality is small; paired projection rejected + +Two benchmark-only CUDA prototypes were completed against the production +native NVFP4 grouped kernels at source `a5bceb3006`. Both use actual layer-0, +TP4-rank-0 checkpoint weights, synthetic FP16 activations, captured C16 +routes (sliced for M4/M8), distinct-expert and shared-ten-expert patterns. +Five alternating-order timing samples use 16 calls per CUDA Graph and ten +replays per sample. Complete MoE includes the route planner, W13, W2 and +original-slot-order FP32 weighted reduction. No production dispatch changed. + +W2 locality maps four warps to `(4 groups, 1 output tile)`, `(2 groups, +2 tiles)` or `(1 group, 4 tiles)`. For the captured 99-group C16 case, +complete MoE changes **125.978 -> 123.782 us** with two adjacent N tiles; +W2-only changes **42.099 -> 40.083 us**. The four-tile complete result is +123.795 us. The isomorphic one-tile sidecar already measures 124.525 us, +so the entire original-to-four-tile gain cannot be attributed to locality: +all sidecar variants use 52 registers versus the original's 44. At M4, +distinct/captured complete timings with four tiles regress +58.758/59.283 -> 59.488/59.667 us. M4 is a grouped-kernel screen, not the +current production M4 dispatch. Seven changed-input/route/poisoned-buffer +replays produce exactly equal final outputs at every screened case. +The layer-0 C16 saving projected over 48 calls is only about **0.105 ms**, +not a measured model gain. Do not run an endpoint battery or enable it by +default on that basis. + +The W13 paired-projection prototype reuses one input fragment for gate and +up in the same warp, retaining separate accumulators, original Split-K sum +order, and FP16 projection/SiLU/product boundaries. It passes seven changed +graph-replay cases with exactly equal W13 intermediates and final outputs, +but is slower. Captured M4/M8/M16 complete MoE: +**58.842/90.554/125.894 -> 62.099/99.584/129.677 us**. C16 W13 alone: +**78.874 -> 83.296 us**. Distinct and shared-ten patterns also regress. +Reject this implementation. For Split8, registers grow 40 -> 63 and CTA +threads shrink 512 -> 256, with the same 17408 shared bytes and no spills. +Static resource bounds allow three original CTAs (48 warps) versus four +paired CTAs (32 warps) per SM. These are compiler/occupancy API limits, not +measured active warps, achieved bandwidth or proof of a particular stall. + +Sources: `benchmark_sm70_moe_w2_locality.py`, `sm70_moe_w2_locality.cu`, +`sm70_moe_w13_paired.cu` under `benchmarks/kernels/`. CPU-only builds use +CUDA 12.8, Torch 2.10.0+cu128, explicit `TORCH_CUDA_ARCH_LIST=7.0` and no +CUDA device initialization. W2/paired binary SHA256 respectively: +`95ddcccad85149d873611a37cfcd7f2667e4a3382ccee9bf7294c8806efd8fdb`, +`4c3c0c0cfb8cbf86f789ee4ac69793d64623d158d868ab9f2e5e3f93de01c5f0`. +Artifacts: `.artifacts/moe-{w2-locality,w13-paired}/screen-v1.{json,log,env,source-sha}`. +Both GPU jobs exited; no model was loaded and no API remains from them. + +This follows the locality/Split-K questions in the +[PyTorch MoE kernel report](https://pytorch.org/blog/accelerating-moe-model/), +not its A100/H100 performance claims. The +[Volta tuning guide](https://docs.nvidia.com/cuda/archive/11.0/volta-tuning-guide/index.html) +informs the static resource interpretation; it does not establish our +achieved occupancy. + +The old 48-layer C16 route capture contains 87.292 unique experts and +87.854 eight-row groups per layer on average. Each valid group logically +loads 409600 weight + 102400 scale bytes for W13 and 204800 weight + 51200 +scale bytes for W2. Total logical weight/scale traffic is **3,238,656,000 +bytes**. Dividing by the V100's nominal 900 GB/s gives 3.599 ms, but this +is only an ideal streaming estimate: the cohort is not the current Nsight +sample, caches can reuse group data, and instructions/activations/outputs +are not included. Do not present the ratio to 6.85 ms as measured DRAM +utilization or a guaranteed removable wall cost. The remaining full-step +gap will require more than the small W2 mapping gain. + +Source/stream ordering also shows that shared-expert projection, sigmoid +and multiply run on the stream opposite routed MoE. For example, one +interior graph's first layer finishes the shared branch at 0.494 ms while +W2 finishes at 0.552 ms. Its inferred ~1.0-ms total service is not a +separate 1.0-ms critical-path saving. Do not sum a prospective shared-gate +fusion gain with MoE service as if both were independent endpoint wins. + +### Scale-layout screen and a profitable copy-only GDN batch fallback + +The tile-major scale-layout MoE screen has also completed. It transposes +only prepared FP16 scale storage to match the packed weights' traversal; +both matrix products, group planning, reductions and rounding remain +unchanged. All seven changed-input/route poisoned replays match the original +W13 intermediate and complete output. At captured C16, complete MoE is +**125.408 -> 122.438 us** with both layouts changed; W13-only layout is +122.880 us. At distinct M4, changing both layouts instead regresses +**57.530 -> 58.592 us** (1.85%). W13-only has no >1% regression in this +small screen but its saving is too small to justify a production repack or +full-model launch. Do not count the ~0.14-ms layer-0 projection as an +endpoint gain or claim zero extra production scale storage. Keep this +variant benchmark-only. Binary SHA256: +`1a3394849e2739e7b4a1bafb53dcbf52ddf31ccbdeacec4a87bf342888ceb1f4`. +Artifacts: `.artifacts/moe-scale-layout/screen-v1.{json,log,env,source-sha}`; +the compiled source is retained in `.artifacts/sm70_moe_scale_layout.cu` +and its formatted benchmark copy. The job respected live foreign ownership, +then completed and released GPU0. + +Source/trace audit found a separate, avoidable serial batch cost: the opaque +single-token GDN input custom op falls back at M>1 to **two original GEMMs +plus four `.contiguous()` kernels**. The four copies are visible immediately +after GDN b/a reduction on all 36 GDN layers. The new copy-only kernel +combines those copies into one; it does not concatenate GEMM weights, +reassociate accumulation, fuse nonlinear arithmetic, alter scales or use +INT8. The original single-token GEMV remains unchanged. + +The standalone micro uses real layer-0/TP4-rank-0 FP16 weights, synthetic +inputs and 16 distinct allocations of the same weights (not 16 actual +layers), exceeding L2. Five alternating graph samples include both unchanged +GEMMs and all output copies. After moving the kernel into the runtime module, +the same screen rechecks that exact source: + +| M | Original complete input projection, us | Fused copies, us | +| ---: | ---: | ---: | +| 2 | 62.150 | 51.088 | +| 4 | 59.430 | 48.950 | +| 8 | 59.504 | 49.427 | +| 16 | 60.774 | 50.509 | +| 32 | 62.077 | 52.250 | +| 64 | 80.662 | 71.581 | + +C16 copy-only is approximately 7.0 -> 1.5 us. The complete-chain saving +projects to about **0.37 ms over 36 layers**, not a measured engine gain. +This will not independently close the remaining 5.15-ms C16 step gap. +All 65536 FP16 bit patterns, including NaN payloads, pass each store branch; +four poisoned changed-input graph replays at each production-width screen +are bitwise equal. The public registered op's CPU/GPU suite reports +**35 passed**, including changed input/weights, graph replay, explicit +fusion route-hit tracking, unsupported fallback and unchanged M1. +These are operator gates, not dataset-quality admission. + +Combined CPU regression: **74 passed, eight GPU cases skipped, three old +HC GPU cases explicitly deselected**. The first CPU-only attempt selected +those three HC cases because their existing physical-device check ignores +`CUDA_VISIBLE_DEVICES=''`; all failed during CUDA initialization, not in an +operator. The corrected invocation excludes only those three GPU tests. +The new suite's eight CUDA cases are covered by the owned 35-pass GPU run. + +Runtime opt-in: `VLLM_SM70_GDN_BATCH_SPLIT_COPY=1`, **default 0**. Local +admission requires SM70, packed 2D FP16 QKVZ/b/a outputs with the existing +custom op's local dimensions and M>1. It has no upper batch, maxseq, chunk, +KV dtype or checkpoint-name condition. CPU, empty, M1, strided, wrong-dtype +and unsupported outputs retain the original fallback. No persistent cache +or extra model-weight copy is introduced. The existing custom op's fake +output shapes and non-aliasing/contiguity contract stay unchanged. + +Micro/runtime artifacts: +`.artifacts/gdn-projection-split/screen-v1.{json,log}` and +`.artifacts/gdn-projection-split-runtime/{screen-v2.json,screen-v2.log,pytest-v2.log}`. +Both jobs exited and released their GPU processes. No production default +was enabled. Next combine this local gain with the existing batch candidate +only after matching endpoint and original-production quality comparisons; +the 238/420/728 tok/s targets and full quality admission remain open. + +### Original-production API comparator completed: quality gate remains open + +The missing original-production control now completed at source +`eb40e649ac02fddfb0652a83bd96a369ba69785d`, using the same pinned native +owner and Flash-V100 package. All existing runtime settings match the +retained candidate API workload except the experimental grouped MoE and +batch HC flags are explicitly **0**. The new GDN split-copy flag is also +explicitly **0**; the retained candidate predates that implementation. +Thus this comparison does **not** test the new copy-only optimization. +The idle profiler wrapper/configuration is omitted; no new speed trace or +throughput result is claimed. + +All 80 request payloads, seeds, dataset hashes and order match the retained +candidate, verified before and after execution. Sampling remains temperature +1.0/top-k20/top-p0.95, thinking off, max16384 and natural termination. +Client peak is 16; mixed request lengths are not a fixed GPU C16 benchmark. +All requests succeed with 50 `tool_calls` and 30 `stop`, no truncation. + +| Subset | Original production | Earlier grouped-MoE + batch-HC candidate | +| --- | ---: | ---: | +| BFCL simple | 15/16 | 14/16 | +| BFCL parallel | 12/16 | 12/16 | +| BFCL multiple | 14/16 | 14/16 | +| BFCL irrelevance | 13/16 | 11/16 | +| **BFCL total** | **54/64** | **51/64** | +| JSONSchemaBench | **16/16** | **16/16** | + +There are four candidate losses and one candidate win: `simple_python_7` +(unit spelling outside the reference alternatives), `parallel_0` (one of +two requested calls omitted), `irrelevance_4` and `irrelevance_8` (unneeded +tools called) versus the candidate win on `parallel_8` (control emits the +wrong qualified tool name). This is an adverse small-sample task-score +signal, **not established causal degradation**. The one-sided exact paired +sign probability for at least four losses among five discordant cases is +6/32 = 0.1875; it neither establishes a regression nor proves noninferiority. +Do not promote existing batch HC/grouped-MoE numerics from the candidate- +only scores or use the new copy kernel's bitwise checks to dismiss it. + +Next localize the signal by separating grouped-MoE and HC in a matched +diagnostic ablation, preserving the full fixed quality manifest for admission. +Do not alter the accepted answer list or report an easier discordant-only +subset as a recovered dataset score. Endpoint/C1 and additional coding/ +teacher-forced-decode quality gates remain outstanding; all production +defaults remain unchanged and PR #504 remains Draft. + +Artifacts: `.artifacts/batch-original-api-quality/{control-tools-v1.json, +paired-v1.json,control-v1.env,server-v1.log,client-v1.log}` and +`.artifacts/run_original_api_quality.sh`. Raw control response SHA256: +`54c971cc863b1d4f7ec5ffc3ef95984aea1ef187294f22814ef04d4f0b6fdfcf`. +The finite launcher reports `client_status=0 server_status=0`. API PID +2030560 and worker PIDs 2030652--2030655 exited; GPU0--3 have no owned +compute process after cleanup. Multiprocessing printed semaphore/shared- +memory tracker shutdown warnings; no persistent GPU allocation was observed. +No API or queued GPU job remains from this task. + +### Separate MoE/HC quality arms and fixed-trajectory decode diagnosis + +At the same runtime source `b6c18db202`, the two missing factor arms ran the +unchanged 80-case API manifest: grouped MoE only and batch HC only. All +payloads, prompt token IDs, seeds and dataset hashes match earlier runs. +Both return 80 successful natural completions, with no JSON Schema failures. +Worker logs confirm grouped MoE at graph widths 8/16 in the first arm and +batch HC at widths 2/4/8/16 in the second. The GDN copy flag remains 0. + +| Natural generation subset | Original | MoE only | HC only | Both | +| --- | ---: | ---: | ---: | ---: | +| Simple | 15/16 | 14/16 | 14/16 | 14/16 | +| Parallel | 12/16 | 11/16 | 12/16 | 12/16 | +| Multiple | 14/16 | 14/16 | 14/16 | 14/16 | +| Irrelevance | 13/16 | 12/16 | 12/16 | 11/16 | +| **BFCL total** | **54/64** | **51/64** | **52/64** | **51/64** | +| Schema | 16/16 | 16/16 | 16/16 | 16/16 | + +The loss sets differ. Both isolated arms miss `simple_python_14`, +`parallel_1`, `irrelevance_1`, and fix control's `parallel_8`; MoE-only +additionally misses `parallel_15`. These are not the earlier combined arm's +four losses. Inspection finds a JSON code block instead of a protocol call, +missing calls, or an unnecessary tool, not transport errors silently dropped +by the scorer. No answer alternatives or scoring rules were changed. +All prompt IDs match; 31 outputs differ from original in each arm, including +five first-token differences for each isolated arm and eight for the earlier +combined arm. First-token differences cannot be attributed solely to a +decode-only arithmetic change without examining prefill/initial sampling. + +The [upstream reproducibility documentation](https://docs.vllm.ai/en/stable/usage/reproducibility/) +does not promise online reproducibility from seed alone. Local MRv2 source +keys its Gumbel noise by request seed, absolute position and token ID, not +batch row, so the observation is not itself proof of a row-dependent RNG +bug. We do not enable batch-invariant mode or change production scheduling +to obtain a passing score; upstream's documented batch-invariance hardware +requirement is SM80+, and this would change the performance contract. + +To separate free-running branch changes from conditional model probabilities, +an **artifact-only diagnostic sampler** retains 64 original BFCL trajectories, +up to 64 tokens each (3654 total). Reserved seeds 30260905+index force only +these trajectories; ordinary 20260905+index quality requests pass through the +original sampler untouched. Standard MRv2 raw-logprob reporting then scores +the forced chosen token before temperature/top-k/top-p. These are model- +generated control continuations, **not human references, a corpus PPL score, +or free-running quality acceptance**. They must never be deployed or counted +as throughput. CPU mapping checks cover unknown seeds, partial prefill, +position limits and row permutation. Every API response must match the full +reference token sequence and return finite logprobs. + +One additional unchanged natural-quality replay precedes the diagnostic in +each server. Original now scores **55/64**, combined candidate **54/64**; +both still score Schema 16/16. Retain these alongside the earlier 54/51, +not as replacements or a best-of selection. Same-seed online scores move; +the available samples still do not establish quality noninferiority. + +All 3654 forced tokens match in both arms. Per-worker telemetry verifies +GPU/CPU seed copies and identical row/position coverage across all four TP +ranks. Control includes 116 pure C16 decode steps; the comparison uses only +matching case/token positions when reporting same-width results: + +| Diagnostic scope | Tokens | Mean candidate-minus-control NLL | Mean absolute delta | Max absolute delta | +| --- | ---: | ---: | ---: | ---: | +| First token / prefill | 64 | -0.0432653 | 0.1700534 | 1.6726222 | +| All continuation decode | 3590 | -0.0009651 | 0.0126083 | 2.6824248 | +| Matched pure C16 decode | 1622 | +0.0013119 | 0.0084666 | 2.6824248 | + +Positive NLL delta means lower candidate probability for the retained token. +The per-case mean decode delta's paired bootstrap 95% interval is +[-0.0031603, +0.0015295] nats; it does not show a clear average deterioration +or prove noninferiority. The large tails must not be hidden by that average. +The largest same-width tail is `parallel_7`, offset27, token271, prompt length +301: control logprob -0.657897 versus candidate -3.340321. Mixed-width rows +and prefill-state differences remain confounders even with fixed continuations. +No C4 same-width intersection exists and C8 has only 19 tokens: do not call +this complete C4/C8 quality admission. + +Relevant existing repair found during source review: +[PR #494](https://github.com/1CatAI/1Cat-vLLM/pull/494), fixed SHA +`5fa8a605dab12cc9ee15459d9ac6b88d95c7be3a`, stabilizes physical-KV-allocation- +dependent page4 attention reduction order. Its author demonstrates a causal +prefill defect on AWQ and also documents independent HC reduced-precision +GEMM and FP16 collective effects. **That does not establish the cause of +our NVFP4 tails.** Next validate this existing narrow repair rather than +duplicate it. Its branch's complete `qsa.py` is older: replacing that whole +file would remove our batched two-warp partial, output-gate fusion and +selector sidecar support. Only the PR's planner delta may be reused; a +rebuilt Flash-V100 binary is required. No PR was merged and no runtime fix +or precision/NCCL default from that investigation has been applied here. + +Artifacts: `.artifacts/batch-api-{moe-only,hc-only}/ablation-tools-v1.json`, +`.artifacts/batch-teacher-{control,candidate}/` (natural responses, forced +responses and four `teacher-rows-.jsonl` files), +`.artifacts/batch-teacher-comparison-v1.json`. Manifest SHA256: +`cb7afc127b1909399f811b216e288b680077695a651c55ef3e842596fff44461`. +Forced control/candidate response SHA256: +`5a479f76d0ebf4da4a0b1698c749d46f703042326cffb9b49657ef2f1ecc5906` / +`8c8171226767ddfd2696db9123b724911c8a81b149e2e137e7e8d634289b6794`. +The four finite API jobs completed with client/server status 0; all their +workers exited and ports 18186/18187/18189/18190 closed. No owned GPU model +or service remains. An initial teacher launcher syntax error was caught by +`bash -n` before any GPU launch and repaired in the artifact script. + +### One physical N32 tile per MoE CTA: reject both resource variants + +Interleaved W13 gate/up pairs fit within a single physical N32 tile. The new +benchmark splits the old two-tile CTA into two independent CTAs, keeping +identical dot products, Split-K accumulation order and FP16 boundaries. +All seven changed-route/input/poisoned graph cases preserve intermediate +and final outputs. No production dispatch changes. + +The unbounded candidate is slower on the captured route at M4/M8/M16: +complete MoE **58.886/90.502/125.811 -> 62.989/91.322/128.250 us**. +At C16, registers grow 40 -> 52 despite halving shared memory 17408 -> 8704 +bytes and CTA threads 512 -> 256. Static occupancy limits drop 48 -> 32 warps. +The second, compile-time `MOE_SINGLE_TILE_BOUND` screen constrains Split8 +to six resident CTAs: compiler resources become 40 registers, 8704 shared +bytes and zero spills, with a 48-warp resource ceiling. It is still slower: +captured C16 complete MoE **126.093 ->139.232 us**, W13 alone +**78.547 ->90.957 us**. Better theoretical occupancy is not proof of better +actual throughput. No NCU achieved-occupancy/bandwidth claim is made. +Both variants help some shared-ten-expert synthetic cases but fail the real- +route gate; reject them without endpoint reruns or a workload-specific +production switch. The bounded follow-up stops at C16 instead of expanding +a failing variant across more widths. + +Sources: `benchmarks/kernels/sm70_moe_w13_single_tile.cu` and the existing +paired screen's `--w13-single-tile`. Binary unbounded SHA256: +`eeff87ee7846240295f8557de6ce137aa351a9fef52f7342359dd31eb3ecd162`. +Artifacts: `.artifacts/moe-w13-single-tile{,-bound}/screen-v1.{json,log}`, +CPU-only build logs and finite task launchers. Both GPU jobs have exited. +The last actual C4/C8/C16 throughput is still 216.496/364.129/589.860 tok/s, +not new measurements from these failed micros. Goals remain unmet. + +### Validate the existing QSA page-order repair without replacing batch code + +At source `46f2b79022`, validate existing +[PR #494](https://github.com/1CatAI/1Cat-vLLM/pull/494), SHA +`5fa8a605dab12cc9ee15459d9ac6b88d95c7be3a`, as an isolated dependency. +The unchanged Flash-V100 source matches that PR's parent; only its CUDA +planner delta (77 additions / 25 deletions) and Python planner delta +(8 additions / 7 deletions) are applied in +`.artifacts/pr494-validation/`. No whole-file replacement, merge, production +route change, or duplicate fix PR. Current batched two-warp, output-gate and +selector-sidecar code remains intact. + +Build: the recorded Python 3.12/Torch 2.10 + CUDA 12.8 environment, SM70, +`MAX_JOBS=4`, CPU-only `setup.py build_ext`, private build-lib/build-temp. +New Flash-V100 SO SHA256: +`b439320b4cd67c0c1c59d41277401a32472770c9995cc82f3b3eb604edb47434`. +An exact-module import hook selects only the modified QSA file and pinned FA +package; the model, other native operators and sidecar stay frozen. + +Run the PR's 26 regressions, current tree's 21 QSA tests, and two additional +8K/page784 contiguous/interleaved checks. Old control: **25 failed, 24 passed**; +fixed: **49 passed**. Failures include logical plans and physical-relocation +attention equality, not CUDA import or device-initialization failures. This +proves the allocation-order defect is present in our frozen components and +the repair composes with current QSA changes. It does not prove the entire +NVFP4 model's numerical or task-quality invariance. + +The two fixed-package API arms reuse the full unchanged 80-case manifest, +official sampling, natural EOS and 16K output cap, followed by the same +3654-token forced-reference diagnostic. Both finite servers finish with +client/server status 0. An occupied foreign port 18191 is left untouched; +the owned jobs use 18201/18202 instead. Actual worker logs confirm grouped +prefill and grouped/XQA tails with production page784; no decode threshold +is lowered. + +| Fixed-package natural subset | Original MoE/HC | Grouped MoE + batch HC | +| --- | ---: | ---: | +| Simple | 14/16 | 16/16 | +| Parallel | 12/16 | 11/16 | +| Multiple | 14/16 | 14/16 | +| Irrelevance | 13/16 | 12/16 | +| **BFCL** | **53/64** | **53/64** | +| Schema | 16/16 | 16/16 | + +All 3654 reference tokens match in both arms, including four-worker trace +agreement. Decode mean candidate-minus-control NLL is **-0.0014224**; the +paired per-case bootstrap interval is **[-0.0033934, +0.0002685]**. The +matched pure-C16 subset has 1641 tokens, mean **+0.00005284** and mean +absolute delta **0.0059487**. However, `parallel_3` offset41 still changes +from logprob -0.4567225 to -2.6368234 (**2.1801 nats**, same decode width). +First-token maximum absolute delta is 1.9212 nats. Equal aggregate scores +and small mean errors are not full quality admission. Compared with older +runs, batch histories and prefill packing remain confounders; do not attribute +all changes, or all remaining tails, to #494. No further full-model rerun is +justified solely to seek a better same-seed score. + +Artifacts: `pr494-validation/{build-v1.log,pytest-control.xml,pytest-fixed.xml}`, +`teacher-{control,candidate}/` and `batch-teacher-comparison-v1.json`. +Forced response SHA256 control/candidate: +`e3cc2431a13c2f40ae9efb068f71c5cab2acc65edc5feb939ec67aa45e437a95` / +`cc3d13047fc9a3d26ce0d9a0350a5a43d24d0ef2bd7f2de17b693e78ee68a406`. +This quality comparison does **not** enable the new GDN copy or the following +attention experiments. + +### Sparse-QSA split-KV: localize planner overhead before a native redesign + +The previously unadmitted grouped small-batch screen now passes all six +changed-input/physical-table/causal-tail-residue eager-versus-graph checks. +Its benchmark had a stale positional call: current QSA inserts an optional +output gate before position/length. Change that benchmark call to keywords. +The first attempt fails before timing at output-gate shape validation; no +bad result is counted. The corrected test uses current QSA plus the repaired +FA binary, not the old worktree's Python module. + +Independent requests, 8K context, FP16 KV, page784, Hq/Hkv/D=6/1/256, +16 calls per graph, 40 replays, five alternating samples: + +| Rows | Current Triton | Direct XQA | Grouped, no split | +| --- | ---: | ---: | ---: | +| 4 | 53.232 us | 301.546 us | 1770.187 us | +| 8 | 99.627 us | 306.134 us | 3535.712 us | +| 16 | 166.126 us | 308.522 us | 3690.357 us | + +Source confirms the alternative grouped route launches +`grid=(1,1,num_groups)`: C16 has only **two forward CTAs**. This is a defect +of promoting the prefill-oriented alternative to decode, **not** the route +currently responsible for production C16's 27.125-ms complete step. + +Prototype 1 reuses the existing native arithmetic unchanged: partition the +logical page plan into tile-aligned pieces, replicate Q, run approximately +80 virtual groups, then merge FP16 partial output with FP32 natural-log LSE. +Complete candidate M4/M8/M16 becomes **134.869/195.360/266.766 us**, versus +paired Triton **53.614/99.368/166.803 us**. This recovers parallelism but +still loses to production. No endpoint run or dispatcher change. + +A single C16 phase screen isolates graph service: planner **126.149 us**, +Q/plan packing **7.451 us**, native forward **114.710 us**, merge **6.213 us**. +These separately replayed phase costs are diagnostics, not an additive E2E +closure. The first phase-only harness attempt left inference mode and failed +on an in-place tensor update after the whole-call check; the corrected +inference-mode attempt passes. It is not an operator arithmetic failure. + +Prototype 2 removes the cross-query hash union in this **benchmark only**. +Reuse the existing XQA logical-token-to-physical-page kernel; retain each +query's selection order, give it disjoint mask bits and tile-padded plan +slots, then use the same split/forward/merge adapter. Physical pages shared +by different queries need not be deduplicated to preserve their separate +contributions, but this prototype is **not** accepted for production prefix +sharing, malformed metadata, E4M3 or full-model quality. + +| Rows | Paired Triton | Per-query plan + split | Decision | +| --- | ---: | ---: | --- | +| 4 | 53.574 us | 94.070 us | Reject regression | +| 8 | 99.203 us | 96.672 us | Small initial micro gain only | +| 16 | 167.611 us | 145.221 us | 13.36% local micro gain only | + +All six changed-input/table/tail-residue graph checks pass against eager and +FP32 attention. Maximum relative L2 is 0.000360; maximum absolute error +0.0017114. These operator tolerances do not admit model quality. C16's +22.390-us saving projects to only **0.269 ms across 12 QSA layers**, not a new +complete-step result; C4 fails the <=1% regression gate. Do not hide it by +shipping an exact-C16-only switch. + +The next structural question is a native independent-request GQA tile: +the reused verifier reserves 48 query/head rows and 512 threads even where +only a small subset is active. Inspect existing XQA/GQA helpers before a +smaller native tile, preserve explicit logical plans and FP32 merge, then +micro-screen the full batch family. Avoid further blind planner/split-count +sweeps. Upstream +[vLLM #54873](https://github.com/vllm-project/vllm/pull/54873) valid-count +skipping is relevant to short-context padding, but cannot erase our full +2048-entry 8K selection; do not import its newer-GPU speedup as our evidence. + +Artifacts under `.artifacts/pr494-validation/`: +`sparse-micro-v2.json`, `sparse-split-v1.json`, +`sparse-split-phases-v2.json`, `sparse-split-direct-v1.json`. +Prototype snapshots are `sparse_split_v1.py`, `sparse_split.py` and +`sparse_split_direct.py`; the last SHA256 is +`1ee8e2f1552ede80444b22d606f4ae21416de74f284ff7c563784dee6827ae74`. +The canonical-index CPU suite passes 24 cases after the benchmark call repair. +All finite model/micro jobs exited and owned GPU allocations were released. +The throughput goals are still unmet, and no experimental default is enabled. + ## Acceptance gates - A microbenchmark candidate must improve median CUDA Graph replay time at its diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 7384538f75..89a61dc6cf 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -14,6 +14,8 @@ #include #include #include +#include +#include #include "fp8_kv_utils.cuh" #include "fused_mma.h" @@ -3659,10 +3661,23 @@ constexpr int kGroupedSparseQueries = 8; constexpr int kGroupedSparsePlannerThreads = 512; constexpr int kGroupedSparseHashCapacity = 8192; constexpr unsigned long long kGroupedSparseEmptyEntry = 0x00000000ffffffffULL; +constexpr int kGroupedSparseItemsPerThread = + kGroupedSparseHashCapacity / kGroupedSparsePlannerThreads; +using GroupedSparseSort = + cub::BlockRadixSort; +// Hash entries plus logical owners exactly fit Volta's 96 KiB opt-in limit. +// After loading both into registers, reuse this storage for sorting and scans. +constexpr size_t kGroupedSparsePlannerSharedMemory = + kGroupedSparseHashCapacity * + (sizeof(unsigned long long) + sizeof(uint32_t)); +static_assert(sizeof(GroupedSparseSort::TempStorage) <= + kGroupedSparsePlannerSharedMemory); __device__ __forceinline__ void grouped_sparse_hash_insert( - unsigned long long* __restrict__ hash_table, const int physical_microblock, - const uint32_t token_mask) { + unsigned long long* __restrict__ hash_table, + uint32_t* __restrict__ logical_owners, const int physical_microblock, + const uint32_t token_mask, const int query, const int logical_token) { if (physical_microblock < 0 || token_mask == 0) { return; } @@ -3675,13 +3690,17 @@ __device__ __forceinline__ void grouped_sparse_hash_insert( for (int probe = 0; probe < kGroupedSparseHashCapacity; ++probe) { const unsigned long long old = atomicCAS(hash_table + slot, kGroupedSparseEmptyEntry, desired); - if (old == kGroupedSparseEmptyEntry) { - return; - } - if (static_cast(old) == - static_cast(physical_microblock)) { + if (old == kGroupedSparseEmptyEntry || + static_cast(old) == + static_cast(physical_microblock)) { atomicOr(hash_table + slot, static_cast(token_mask) << 32); + // Nonnegative int32 tokens use at most 29 bits after division by four. + // The first contributing query owns shared pages, independent of request + // slot IDs, physical allocation, insertion order, and hash collisions. + const uint32_t owner = (static_cast(query) << 29) | + (static_cast(logical_token) >> 2); + atomicMin(logical_owners + slot, owner); return; } slot = (slot + 1) & (kGroupedSparseHashCapacity - 1); @@ -3727,6 +3746,44 @@ __device__ __forceinline__ int grouped_sparse_active_m_tiles( return active_m_tiles; } +// Sort only the device-compacted live union. The physical-page union, logical +// owner tie break, category padding and attention arithmetic are unchanged. +// Compact storage, sorting scratch and category scans reuse the same 96 KiB. +template +__device__ __forceinline__ void grouped_sparse_sort_compacted( + unsigned long long* entries_smem, uint32_t* owners_smem, + const int live_entries) { + using Sort = + cub::BlockRadixSort; + static_assert(sizeof(typename Sort::TempStorage) <= + kGroupedSparsePlannerSharedMemory); + unsigned long long entries[ITEMS_PER_THREAD]; + unsigned long long keys[ITEMS_PER_THREAD]; +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + const int index = threadIdx.x * ITEMS_PER_THREAD + item; + const auto entry = + index < live_entries ? entries_smem[index] : kGroupedSparseEmptyEntry; + entries[item] = entry; + const int category = + grouped_sparse_active_m_tiles(static_cast(entry >> 32)); + keys[item] = index < live_entries + ? (static_cast(category) << 32) | + owners_smem[index] + : ULLONG_MAX; + } + __syncthreads(); + auto& scratch = *reinterpret_cast(entries_smem); + Sort(scratch).Sort(keys, entries, 0, 36); + __syncthreads(); +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + entries_smem[threadIdx.x * ITEMS_PER_THREAD + item] = entries[item]; + } + __syncthreads(); +} + __global__ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_plan_kernel( const int* __restrict__ logical_indices, @@ -3741,15 +3798,13 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla const int physical_page_stride, const int num_cache_blocks) { const int group_idx = blockIdx.x; const int tid = threadIdx.x; - __shared__ int category_counts[8]; - __shared__ int category_offsets[8]; - __shared__ int category_cursors[8]; - __shared__ int - warp_category_prefix[(kGroupedSparsePlannerThreads / kWarpSize) * 8]; extern __shared__ unsigned long long hash_table[]; + auto* logical_owners = + reinterpret_cast(hash_table + kGroupedSparseHashCapacity); for (int slot = tid; slot < kGroupedSparseHashCapacity; slot += kGroupedSparsePlannerThreads) { hash_table[slot] = kGroupedSparseEmptyEntry; + logical_owners[slot] = UINT_MAX; } __syncthreads(); @@ -3793,8 +3848,9 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla first_token, request_idx, request_block_table, request_block_table_stride, block_table_width, page_size, physical_page_stride, num_cache_blocks); - grouped_sparse_hash_insert(hash_table, physical_microblock, - 0xFu << (query * 4)); + grouped_sparse_hash_insert(hash_table, logical_owners, + physical_microblock, 0xFu << (query * 4), + query, first_token); } else { #pragma unroll for (int token_offset = 0; token_offset < 4; ++token_offset) { @@ -3804,8 +3860,9 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla token, request_idx, request_block_table, request_block_table_stride, block_table_width, page_size, physical_page_stride, num_cache_blocks); - grouped_sparse_hash_insert(hash_table, physical_microblock, - 1u << (query * 4 + (token & 3))); + grouped_sparse_hash_insert( + hash_table, logical_owners, physical_microblock, + 1u << (query * 4 + (token & 3)), query, token); } } } @@ -3842,19 +3899,85 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla request_block_table_stride, block_table_width, page_size, physical_page_stride, num_cache_blocks); const uint32_t tail_mask = ((1u << tail_count) - 1) << (query * 4); - grouped_sparse_hash_insert(hash_table, physical_microblock, tail_mask); + grouped_sparse_hash_insert(hash_table, logical_owners, + physical_microblock, tail_mask, query, + selected_tail_token); } } } __syncthreads(); + // Keep the existing physical-page union and masks, but never let physical + // hash slots determine the attention reduction order. Category is primary + // to preserve active-tile packing; the logical owner orders each category. + unsigned long long entries[kGroupedSparseItemsPerThread]; + uint32_t owners[kGroupedSparseItemsPerThread]; + int local_count = 0; +#pragma unroll + for (int item = 0; item < kGroupedSparseItemsPerThread; ++item) { + const int slot = tid * kGroupedSparseItemsPerThread + item; + const unsigned long long entry = hash_table[slot]; + entries[item] = entry; + owners[item] = logical_owners[slot]; + local_count += static_cast(entry) != 0xffffffffu; + } + __syncthreads(); + using Scan = cub::BlockScan; + auto& scan_storage = *reinterpret_cast(hash_table); + int compact_offset = 0, live_entries = 0; + Scan(scan_storage).ExclusiveSum(local_count, compact_offset, live_entries); + __syncthreads(); + // Publish the aggregate outside BlockScan scratch, then read it before + // compacted owner writes can overwrite this temporary scalar. + if (tid == 0) logical_owners[0] = live_entries; + __syncthreads(); + live_entries = logical_owners[0]; + __syncthreads(); +#pragma unroll + for (int item = 0; item < kGroupedSparseItemsPerThread; ++item) { + if (static_cast(entries[item]) != 0xffffffffu) { + hash_table[compact_offset] = entries[item]; + logical_owners[compact_offset++] = owners[item]; + } + } + __syncthreads(); + int sorted_items = kGroupedSparsePlannerThreads; + while (sorted_items < live_entries) sorted_items *= 2; + switch (sorted_items / kGroupedSparsePlannerThreads) { + case 1: + grouped_sparse_sort_compacted<1>(hash_table, logical_owners, + live_entries); + break; + case 2: + grouped_sparse_sort_compacted<2>(hash_table, logical_owners, + live_entries); + break; + case 4: + grouped_sparse_sort_compacted<4>(hash_table, logical_owners, + live_entries); + break; + case 8: + grouped_sparse_sort_compacted<8>(hash_table, logical_owners, + live_entries); + break; + default: + grouped_sparse_sort_compacted<16>(hash_table, logical_owners, + live_entries); + break; + } + auto* category_counts = reinterpret_cast(logical_owners); + int* category_offsets = category_counts + 8; + int* category_cursors = category_offsets + 8; + int* warp_category_prefix = category_cursors + 8; + __syncthreads(); + if (tid < 8) { category_counts[tid] = 0; category_offsets[tid] = 0; category_cursors[tid] = 0; } __syncthreads(); - for (int slot = tid; slot < kGroupedSparseHashCapacity; + for (int slot = tid; slot < sorted_items; slot += kGroupedSparsePlannerThreads) { const unsigned long long entry = hash_table[slot]; if (static_cast(entry) != 0xffffffffu) { @@ -3877,7 +4000,7 @@ __launch_bounds__(kGroupedSparsePlannerThreads, 1) void grouped_sparse_page4_pla constexpr int kPlannerWarps = kGroupedSparsePlannerThreads / kWarpSize; const int lane = tid & (kWarpSize - 1); const int warp = tid / kWarpSize; - for (int chunk_start = 0; chunk_start < kGroupedSparseHashCapacity; + for (int chunk_start = 0; chunk_start < sorted_items; chunk_start += kGroupedSparsePlannerThreads) { const unsigned long long entry = hash_table[chunk_start + tid]; const uint32_t physical_microblock = static_cast(entry); @@ -4006,17 +4129,16 @@ at::Tensor flash_attention_grouped_sparse_page4_plan( TORCH_CHECK(properties->major == 7 && properties->minor == 0, "grouped sparse page4 planner supports SM70 only"); cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - constexpr size_t kPlannerSharedMemory = - kGroupedSparseHashCapacity * sizeof(unsigned long long); - const cudaError_t smem_status = cudaFuncSetAttribute( - grouped_sparse_page4_plan_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, kPlannerSharedMemory); + const cudaError_t smem_status = + cudaFuncSetAttribute(grouped_sparse_page4_plan_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + kGroupedSparsePlannerSharedMemory); TORCH_CHECK(smem_status == cudaSuccess, "Failed to set grouped sparse page4 planner shared memory: ", cudaGetErrorString(smem_status)); - grouped_sparse_page4_plan_kernel<<(num_groups), - kGroupedSparsePlannerThreads, - kPlannerSharedMemory, stream>>>( + grouped_sparse_page4_plan_kernel<<< + static_cast(num_groups), kGroupedSparsePlannerThreads, + kGroupedSparsePlannerSharedMemory, stream>>>( logical_indices.data_ptr(), block_table.data_ptr(), token_to_req.data_ptr(), query_positions.data_ptr(), sequence_lengths.data_ptr(), output_blocks.data_ptr(), diff --git a/flashinfer-sm70/README.md b/flashinfer-sm70/README.md index 7061b88b41..e09e0d56a6 100644 --- a/flashinfer-sm70/README.md +++ b/flashinfer-sm70/README.md @@ -41,3 +41,12 @@ reference, requires zero PTXAS spills and zero SASS `LDL`/`STL`, checks for 256 emitted `HMMA.884` instructions per kernel, and reports paired timings. Timing is observational only and is not an attention promotion or speedup gate. + +## Sparse QSA decode port (experimental) + +The separate [QSA source-port worklog](../docs/design/sm70_flashinfer_qsa_port.md) +documents an SM70 instantiation of pinned FlashInfer CUDA decode and cascade +kernels with an ordered sparse-cache adapter. It uses a newer independent +source pin, not the WMMA probe above. It is benchmark-only, has no serving +dispatch/default change, and must pass GPU correctness and measured speed +gates before runtime integration. diff --git a/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_decode.cuh b/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_decode.cuh new file mode 100644 index 0000000000..6728c69851 --- /dev/null +++ b/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_decode.cuh @@ -0,0 +1,147 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// This adapter instantiates the actual upstream CUDA decode and cascade +// kernels. Upstream: +// flashinfer-ai/flashinfer@6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f. No +// Triton/Flash-V100 fallback is hidden behind this interface. +#include + +namespace flashinfer::attention::sm70 { + +// QSA's ordered selection is a virtual page-size-one cache. Keeping raw +// selection order and duplicates is essential: physical-page sorting changed +// numerical behavior in the previous grouped-page4 implementation. +struct UnitPage { + __host__ __device__ operator unsigned int() const { return 1; } + __device__ void divmod(uint32_t n, uint32_t& q, uint32_t& r) const { + q = n; + r = 0; + } +}; + +struct SafeCachePointer { + const half* data; + const half* zero; + __device__ const half* operator+(size_t offset) const { + // Invalid selections must not read arbitrary cache values: 0 * NaN in + // the PV loop is still NaN. Use a persistent zero page, including padding. + return (offset >> 63) ? zero + (offset & 255) : data + offset; + } +}; + +struct SparsePagedKV { + UnitPage page_size; + uint32_t batch_size, num_heads, width; + int64_t head_stride; + SafeCachePointer k_data, v_data; + const int32_t* indptr; + const int32_t* rope_pos_offset = nullptr; + const int64_t* offsets; + + __device__ uint32_t get_length(uint32_t) const { return width; } + __device__ size_t protective_get_kv_offset(uint32_t index, uint32_t head, + uint32_t, uint32_t d, + int32_t end) const { + const int64_t offset = index < end ? offsets[index] : -1; + return offset < 0 ? (size_t{1} << 63) + : size_t(offset + head * head_stride + d); + } +}; + +struct QSAParams { + using DTypeQ = half; + using DTypeKV = half; + using DTypeO = float; + using IdType = int32_t; + const half* q; + float* o; + float* lse; + SparsePagedKV paged_kv; + const bool* block_valid_mask = nullptr; + uint32_t padded_batch_size, num_qo_heads; + bool partition_kv = true; + const int32_t* request_indices; + const int32_t* kv_tile_indices; + const int32_t* kv_chunk_size_ptr; + uint32_t q_stride_n, q_stride_h; +}; + +struct QSAVariant { + static constexpr bool use_softmax = true; + float sm_scale_log2 = math::log2e / 16.f; + __device__ QSAVariant(const QSAParams&, uint32_t, uint8_t*) {} + __device__ float LogitsTransform(const QSAParams&, float value, uint32_t, + uint32_t, uint32_t, uint32_t, + uint32_t) const { + return value; + } + __device__ bool LogitsMask(const QSAParams& p, uint32_t row, uint32_t, + uint32_t index, uint32_t, uint32_t) const { + return index < p.paged_kv.width && + p.paged_kv.offsets[size_t(row) * p.paged_kv.width + index] >= 0; + } + __device__ float OutputTransform(const QSAParams&, float value, uint32_t, + uint32_t, uint32_t, float, float denominator, + float) const { + return denominator > 0.f ? value / denominator : 0.f; + } +}; + +// Same visible-index contract as qsa_sparse_paged_attention. Causality is +// applied by the QSA selector/expander before this operator, not reinterpreted +// as dense causal attention. Invalid rows/pages/indices produce empty states. +__global__ void PrepareQSA(const int32_t* indices, const int32_t* table, + const int32_t* requests, int64_t* offsets, + int32_t* metadata, int rows, int selected, int width, + int splits, int page, int table_width, int nrequests, + int nblocks, int64_t index_stride, + int64_t table_stride, int64_t block_stride, + int64_t token_stride) { + const int row = blockIdx.x; + const int request = requests[row]; + for (int col = threadIdx.x; col < width; col += blockDim.x) { + int64_t offset = -1; + const int logical = col < selected ? indices[row * index_stride + col] : -1; + if (request >= 0 && request < nrequests && logical >= 0 && + logical / page < table_width) { + const int physical = table[request * table_stride + logical / page]; + if (physical >= 0 && physical < nblocks) + offset = physical * block_stride + (logical % page) * token_stride; + } + offsets[size_t(row) * width + col] = offset; + } + if (threadIdx.x == 0) { + metadata[row] = row * width; + if (row == 0) { + metadata[rows] = rows * width; + metadata[rows + 1 + 2 * rows * splits] = width / splits; + } + } + for (int split = threadIdx.x; split < splits; split += blockDim.x) { + metadata[rows + 1 + row * splits + split] = row; + metadata[rows + 1 + rows * splits + row * splits + split] = split; + } +} + +template +void LaunchQSADecode(QSAParams p, half* output, int rows, int splits, + cudaStream_t stream) { + // First port: upstream SIMT/GQA kernel, no tensor-core emulation. On SM70 + // upstream cp_async.cuh uses synchronous vector loads + block barriers. + constexpr int stages = 2, tile = 1, vec = 8, bdx = 32, bdz = 1; + constexpr int smem = 2 * stages * tile * Group * 256 * sizeof(half) + + tile * Group * bdx * sizeof(size_t); + BatchDecodeWithPagedKVCacheKernel + <<>>(p); + // Keep FP32 partials until the final output cast. This is upstream's actual + // cascade implementation, not the previous Triton merge under a new name. + MergeStatesKernel<8, float, half> + <<>>( + p.o, p.lse, output, nullptr, splits, p.num_qo_heads, 256); +} + +} // namespace flashinfer::attention::sm70 diff --git a/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_mqa.cuh b/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_mqa.cuh new file mode 100644 index 0000000000..9f1a7e0981 --- /dev/null +++ b/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_mqa.cuh @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Schedule algorithm adapted from FlashInfer (Apache-2.0), copyright 2025 +// FlashInfer team, and TensorRT-LLM, copyright 2026 NVIDIA CORPORATION. +// Pinned source: 6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f, +// flashinfer/attn_scores/kernels/schedule_kernel.py. +// Native Volta implementation: no CuTe, TMA, FP8 MMA, host length readback, +// global spin barrier, or request-capacity-sized launch grid. +#pragma once + +#include +#include +#include +#include + +#include + +namespace flashinfer::attention::sm70 { + +constexpr int kMQATile = 64; +constexpr int kMQAMaxRows = 64; + +struct MQAParams { + const half* q; + const half* k; + const int32_t* table; + const int32_t* requests; + const void* positions; + const int32_t* lengths; + int32_t* visible; + int32_t* schedule; + float* logits; + int32_t* task_visits; + int rows, heads, columns, pages, page_size, table_width, num_requests; + int ratio, workers; + bool positions64; + float divisor; + int64_t q_row, q_head, k_page, k_token, table_row, out_row; +}; + +// Recomputed on every call/replay into instance-owned storage. In particular, +// empty rows produce duplicate prefixes: upper_bound (<=), not lower_bound. +__global__ void PlanMQA(MQAParams p) { + __shared__ int prefix[kMQAMaxRows]; + const int lane = threadIdx.x; + int carry = 0; + for (int base = 0; base < kMQAMaxRows; base += 32) { + const int row = base + lane; + int visible = 0; + if (row < p.rows) { + const int request = p.requests[row]; + const int64_t pos = p.positions64 + ? static_cast(p.positions)[row] + : static_cast(p.positions)[row]; + if (request >= 0 && request < p.num_requests && pos >= 0) { + // floor(min(pos+1, len)/ratio) equals the minimum of both floors. + // Clamp before adding one: INT64_MAX positions cannot overflow, and + // the bounded covered-token count needs only a 32-bit division. + const int length = max(0, p.lengths[request]); + const int covered = int(min(pos, int64_t(length) - 1) + 1); + const int64_t capacity = int64_t(p.page_size) * p.table_width; + visible = int( + min(int64_t(covered / p.ratio), min(int64_t(p.columns), capacity))); + } + p.visible[row] = visible; + } + int count = (visible + kMQATile - 1) / kMQATile; +#pragma unroll + for (int offset = 1; offset < 32; offset *= 2) { + const int other = __shfl_up_sync(0xffffffff, count, offset); + if (lane >= offset) count += other; + } + count += carry; + prefix[row] = count; + carry = __shfl_sync(0xffffffff, count, 31); + } + __syncwarp(); + const int per_worker = carry / p.workers; + const int extra = carry % p.workers; + for (int worker = lane; worker <= p.workers; worker += 32) { + const int start = worker * per_worker + min(worker, extra); + int lo = 0, hi = p.rows; + while (lo < hi) { + const int mid = (lo + hi) / 2; + if (prefix[mid] <= start) + lo = mid + 1; + else + hi = mid; + } + p.schedule[2 * worker] = lo; + p.schedule[2 * worker + 1] = start - (lo ? prefix[lo - 1] : 0); + } +} + +// QSA is sum_h relu(dot(K, Q_h)) / divisor, not softmax attention. +// FP16 operands, FP32 tensor-core accumulation and a padded 16-head reduction. +// One CTA consumes a balanced contiguous range of *live* (row, KV tile) tasks. +template +__global__ void ScoreMQA(MQAParams p) { + using namespace nvcuda; + // Eight-half skew keeps WMMA's shared-memory rows out of the same banks. + constexpr int LD = Dim + 8; + __shared__ __align__(32) half query[16 * LD]; + __shared__ __align__(32) half keys[kMQATile * LD]; + __shared__ __align__(32) float scores[kMQATile * 16]; + __shared__ int64_t key_offsets[kMQATile]; + const int tid = threadIdx.x, warp = tid / 32; + int row = p.schedule[2 * blockIdx.x]; + int tile = p.schedule[2 * blockIdx.x + 1]; + const int end_row = p.schedule[2 * (blockIdx.x + 1)]; + const int end_tile = p.schedule[2 * (blockIdx.x + 1) + 1]; + while (row < p.rows && + (row < end_row || (row == end_row && tile < end_tile))) { + const int visible = p.visible[row]; + const int tile_end = + row == end_row ? end_tile : (visible + kMQATile - 1) / kMQATile; + if (tile < tile_end) { + for (int i = tid; i < 16 * (Dim / 8); i += 128) { + const int h = i / (Dim / 8), d = (i % (Dim / 8)) * 8; + const int4 value = + h < p.heads ? *reinterpret_cast( + p.q + int64_t(row) * p.q_row + h * p.q_head + d) + : make_int4(0, 0, 0, 0); + *reinterpret_cast(query + h * LD + d) = value; + } + const int request = p.requests[row]; + __syncthreads(); + for (; tile < tile_end; ++tile) { + if constexpr (Audit) { + if (tid == 0) + atomicAdd(p.task_visits + + row * ((p.columns + kMQATile - 1) / kMQATile) + tile, + 1); + } + if (tid < kMQATile) { + const int col = tile * kMQATile + tid; + int physical = -1; + if (col < visible) + physical = + p.table[int64_t(request) * p.table_row + col / p.page_size]; + key_offsets[tid] = physical >= 0 && physical < p.pages + ? int64_t(physical) * p.k_page + + (col % p.page_size) * p.k_token + : -1; + } + __syncthreads(); + for (int i = tid; i < kMQATile * (Dim / 8); i += 128) { + const int local_col = i / (Dim / 8), d = (i % (Dim / 8)) * 8; + const int64_t offset = key_offsets[local_col]; + const int4 value = + offset >= 0 ? *reinterpret_cast(p.k + offset + d) + : make_int4(0, 0, 0, 0); + *reinterpret_cast(keys + local_col * LD + d) = value; + } + __syncthreads(); + wmma::fragment a; + wmma::fragment b; + wmma::fragment c; + wmma::fill_fragment(c, 0.f); +#pragma unroll + for (int d = 0; d < Dim; d += 16) { + wmma::load_matrix_sync(a, keys + warp * 16 * LD + d, LD); + wmma::load_matrix_sync(b, query + d, LD); + wmma::mma_sync(c, a, b, c); + } + wmma::store_matrix_sync(scores + warp * 16 * 16, c, 16, + wmma::mem_row_major); + __syncthreads(); + if (tid < kMQATile) { + const int col = tile * kMQATile + tid; + if (col < visible) { + float values[16]; +#pragma unroll + for (int h = 0; h < 16; ++h) + values[h] = h < p.heads ? fmaxf(scores[tid * 16 + h], 0.f) : 0.f; +#pragma unroll + for (int step = 8; step; step /= 2) +#pragma unroll + for (int h = 0; h < step; ++h) values[h] += values[h + step]; + p.logits[int64_t(row) * p.out_row + col] = + key_offsets[tid] >= 0 ? values[0] / p.divisor : -CUDART_INF_F; + } + } + __syncthreads(); + } + } + ++row; + tile = 0; + } +} + +} // namespace flashinfer::attention::sm70 diff --git a/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_wmma_decode.cuh b/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_wmma_decode.cuh new file mode 100644 index 0000000000..40a338e2a6 --- /dev/null +++ b/flashinfer-sm70/include/flashinfer/attention/sm70/qsa_wmma_decode.cuh @@ -0,0 +1,167 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +// Numerical-compatibility experiment, not an admitted runtime default. +// Reuse the pinned FlashInfer virtual-page preparation and FP32 cascade. +// The Volta partial kernel preserves production's 16-token tile partition, +// FP16 probabilities before PV, and unrounded FP32 normalization state. +#include +#include +#include + +namespace flashinfer::attention::sm70 { + +template +__global__ void QSAWMMACompatiblePartial(QSAParams p, int selected, + int splits) { + using namespace nvcuda; + constexpr int D = 256, LD = D + 8, LP = 24; + __shared__ __align__(32) half query[16 * LD]; + __shared__ __align__(32) half keys[16 * LD]; + __shared__ __align__(32) half values[16 * LD]; + __shared__ __align__(32) half probabilities[16 * LP]; + __shared__ __align__(32) float scores[16 * 16]; + // Reuse one output tile per warp; never reserve a padded 16 x D result. + __shared__ __align__(32) float output[4 * 16 * 16]; + __shared__ float maximum[16], denominator[16], alpha[16]; + __shared__ int64_t offsets[16]; + const int tid = threadIdx.x, warp = tid / 32; + const int row = blockIdx.x / splits, split = blockIdx.x % splits; + const int kv_head = blockIdx.y; + const int first_head = kv_head * Group; + for (int i = tid; i < 16 * (D / 8); i += 128) { + const int h = i / (D / 8), d = i % (D / 8) * 8; + const int4 q = h < Group ? *reinterpret_cast( + p.q + int64_t(row) * p.q_stride_n + + (first_head + h) * p.q_stride_h + d) + : make_int4(0, 0, 0, 0); + *reinterpret_cast(query + h * LD + d) = q; + } + if (tid < 16) { + maximum[tid] = -1e20f; + denominator[tid] = 0.f; + } + // Discover the accumulator fragment's row layout using the supported WMMA + // load API, rather than depending on undocumented lane/register mappings. + for (int i = tid; i < 256; i += 128) scores[i] = float(i / 16); + __syncthreads(); + using Acc = wmma::fragment; + Acc row_map, accum[4]; + wmma::load_matrix_sync(row_map, scores, 16, wmma::mem_row_major); +#pragma unroll + for (int n = 0; n < 4; ++n) wmma::fill_fragment(accum[n], 0.f); + __syncthreads(); + const int tiles = (selected + 15) / 16; + const int start = split * tiles / splits; + const int end = (split + 1) * tiles / splits; + for (int tile = start; tile < end; ++tile) { + if (tid < 16) { + const int index = tile * 16 + tid; + offsets[tid] = + index < selected + ? p.paged_kv.offsets[size_t(row) * p.paged_kv.width + index] + : -1; + } + __syncthreads(); + for (int i = tid; i < 16 * (D / 8); i += 128) { + const int n = i / (D / 8), d = i % (D / 8) * 8; + const int64_t offset = offsets[n]; + int4 k = make_int4(0, 0, 0, 0), v = k; + if (offset >= 0) { + const int64_t address = offset + kv_head * p.paged_kv.head_stride + d; + k = *reinterpret_cast(p.paged_kv.k_data.data + address); + v = *reinterpret_cast(p.paged_kv.v_data.data + address); + } + *reinterpret_cast(keys + n * LD + d) = k; + *reinterpret_cast(values + n * LD + d) = v; + } + __syncthreads(); + if (warp == 0) { + wmma::fragment a; + wmma::fragment b; + Acc c; + wmma::fill_fragment(c, 0.f); +#pragma unroll + for (int d = 0; d < D; d += 16) { + wmma::load_matrix_sync(a, query + d, LD); + wmma::load_matrix_sync(b, keys + d, LD); + wmma::mma_sync(c, a, b, c); + } + wmma::store_matrix_sync(scores, c, 16, wmma::mem_row_major); + } + __syncthreads(); + if (tid < 16) { + float s[16], ps[16]; + float next_max = maximum[tid]; +#pragma unroll + for (int n = 0; n < 16; ++n) { + s[n] = offsets[n] >= 0 + ? scores[tid * 16 + n] * (1.4426950408889634f / 16.f) + : -1e20f; + next_max = fmaxf(next_max, s[n]); + } + const float scale = exp2f(maximum[tid] - next_max); +#pragma unroll + for (int n = 0; n < 16; ++n) { + ps[n] = offsets[n] >= 0 ? exp2f(s[n] - next_max) : 0.f; + probabilities[tid * LP + n] = __float2half_rn(ps[n]); + } +#pragma unroll + for (int step = 8; step > 0; step /= 2) +#pragma unroll + for (int n = 0; n < step; ++n) ps[n] += ps[n + step]; + denominator[tid] = denominator[tid] * scale + ps[0]; + maximum[tid] = next_max; + alpha[tid] = scale; + } + __syncthreads(); + wmma::fragment a; + wmma::fragment b; + wmma::load_matrix_sync(a, probabilities, LP); +#pragma unroll + for (int n = 0; n < 4; ++n) { +#pragma unroll + for (int e = 0; e < Acc::num_elements; ++e) + accum[n].x[e] *= alpha[int(row_map.x[e])]; + wmma::load_matrix_sync(b, values + warp * 64 + n * 16, LD); + wmma::mma_sync(accum[n], a, b, accum[n]); + } + __syncthreads(); + } +#pragma unroll + for (int n = 0; n < 4; ++n) { +#pragma unroll + for (int e = 0; e < Acc::num_elements; ++e) { + const float norm = denominator[int(row_map.x[e])]; + accum[n].x[e] = norm > 0.f ? accum[n].x[e] / norm : 0.f; + } + wmma::store_matrix_sync(output + warp * 256, accum[n], 16, + wmma::mem_row_major); + __syncwarp(); + for (int i = tid % 32; i < Group * 16; i += 32) { + const int h = i / 16, d = warp * 64 + n * 16 + i % 16; + p.o[(size_t(blockIdx.x) * p.num_qo_heads + first_head + h) * D + d] = + output[warp * 256 + i]; + } + __syncwarp(); + } + if (tid < Group) { + const float norm = denominator[tid]; + p.lse[size_t(blockIdx.x) * p.num_qo_heads + first_head + tid] = + norm > 0.f ? maximum[tid] + log2f(norm) : -CUDART_INF_F; + } +} + +template +void LaunchQSAWMMACompatible(QSAParams p, half* output, int rows, int splits, + int selected, cudaStream_t stream) { + QSAWMMACompatiblePartial + <<>>( + p, selected, splits); + MergeStatesKernel<8, float, half> + <<>>( + p.o, p.lse, output, nullptr, splits, p.num_qo_heads, 256); +} + +} // namespace flashinfer::attention::sm70 diff --git a/flashinfer-sm70/include/flashinfer/gdn/sm70/gdn_fused_decode.cuh b/flashinfer-sm70/include/flashinfer/gdn/sm70/gdn_fused_decode.cuh new file mode 100644 index 0000000000..be1749ba05 --- /dev/null +++ b/flashinfer-sm70/include/flashinfer/gdn/sm70/gdn_fused_decode.cuh @@ -0,0 +1,465 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 FlashInfer team +// SM70 adaptation by 1Cat-vLLM contributors. Source: FlashInfer +// 6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f, gdn_fused_decode_sm120.cu. +#pragma once +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#ifndef FI_GDN_IMPL_NAMESPACE + #define FI_GDN_IMPL_NAMESPACE flashinfer::sm70::gdn +#endif +namespace FI_GDN_IMPL_NAMESPACE { + +// Fused GDN decode step for one layer geometry: a single persistent kernel +// covering the in_proj_ba GEMV, the depthwise causal conv1d update (width 4, +// silu), the q/k/v head split, and the gated delta-rule state update with +// qk-L2-norm, replacing the multi-launch serving chain. +// +// The layer geometry is a compile-time parameter of this translation unit, +// supplied by the native wheel translation units or the isolated benchmark. +// Wheel geometries use distinct C++ and Torch namespaces, so several head +// partitions may coexist without template/constant ODR collisions. +// Only the sizes change: the block shape, warp->row mapping and reduction +// trees below are geometry-independent, and the static_asserts state exactly +// which divisibility relations the code relies on. +#if !defined(FI_GDN_HIDDEN) || !defined(FI_GDN_N_BA) || \ + !defined(FI_GDN_QKV_DIM) || !defined(FI_GDN_H_Q) || !defined(FI_GDN_HV) || \ + !defined(FI_GDN_D) || !defined(FI_GDN_CONV_WIDTH) || \ + !defined(FI_GDN_CONV_STATE_LEN) + #error "gdn_fused_decode.cuh requires the FI_GDN_* geometry defines" +#endif +constexpr int HIDDEN = FI_GDN_HIDDEN; +constexpr int N_BA = FI_GDN_N_BA; +constexpr int QKV_DIM = FI_GDN_QKV_DIM; +constexpr int H_Q = FI_GDN_H_Q; +constexpr int HV = FI_GDN_HV; +constexpr int D = FI_GDN_D; +constexpr int CONV_WIDTH = FI_GDN_CONV_WIDTH; +constexpr int CONV_STATE_LEN = FI_GDN_CONV_STATE_LEN; +// v-heads per qk-head: the delta phase maps v-head h to qk-head h/HEADS_PER_QK. +constexpr int HEADS_PER_QK = HV / H_Q; +#ifndef FI_GDN_ROWS_PER_WARP + #define FI_GDN_ROWS_PER_WARP 8 +#endif +constexpr int ROWS_PER_WARP = FI_GDN_ROWS_PER_WARP; +constexpr int GEMV_NSPLIT = 160; +#ifndef FI_GDN_SHARED_PARAMETERS + #define FI_GDN_SHARED_PARAMETERS 0 +#endif +constexpr bool SHARED_PARAMETERS = FI_GDN_SHARED_PARAMETERS; +static_assert(!SHARED_PARAMETERS || D % (8 * ROWS_PER_WARP) == 0, + "a 256-thread CTA must stay within one recurrent head"); +// The gate reduction below unrolls the GEMV partials as 5 warp-wide loads. +static_assert(GEMV_NSPLIT == 5 * 32, + "gate reduction assumes 5 warp-strided loads"); +// The b/a projection produces HV gate values and HV decay values, stored as +// the low and high halves of the N_BA columns (see the delta phase's +// base_b/base_a offsets). +static_assert(N_BA == 2 * HV, "w_ba columns are [b gates | a decays], HV each"); +static_assert(HV % H_Q == 0, + "each qk-head must serve a whole number of v-heads"); +// The delta phase gives each lane 4 consecutive channels of a D-wide row and +// reduces across the warp; the B=1 fast path indexes rows with shifts/masks. +static_assert(D == 4 * 32, + "delta phase maps one D-wide row onto a warp, 4 per lane"); +static_assert((D & (D - 1)) == 0, + "B=1 row index math uses D as a power of two"); +static_assert(D % ROWS_PER_WARP == 0, "warps own whole groups of state rows"); +// mixed_qkv is [q | k | v] with H_Q q-heads, H_Q k-heads and HV v-heads. +static_assert(QKV_DIM == (2 * H_Q + HV) * D, + "qkv_dim must match the head split"); +// The conv phase is unrolled over the width-4 / 3-step shift register below. +static_assert(CONV_WIDTH == 4 && CONV_STATE_LEN == 3, + "conv taps are unrolled as width 4"); + +// log2(D), for the B=1 fast path's shift/mask row indexing. +constexpr int ilog2_ce(int v) { return v <= 1 ? 0 : 1 + ilog2_ce(v >> 1); } +constexpr int D_LOG2 = ilog2_ce(D); + +typedef half f16; + +__device__ __forceinline__ float siluf(float x) { + return x / (1.0f + __expf(-x)); +} +__device__ __forceinline__ float sigmoidf(float x) { + return 1.0f / (1.0f + __expf(-x)); +} +__device__ __forceinline__ float softplusf(float x) { + return x > 20.0f ? x : log1pf(__expf(x)); +} +__device__ __forceinline__ float warp_reduce(float v) { +#pragma unroll + for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o); + return v; +} + +// Cooperative launch guarantees grid residency even when vLLM has auxiliary +// streams. Use CUDA's specified grid synchronization instead of an occupancy- +// capped regular launch with a software spin barrier. +__device__ __forceinline__ void grid_barrier() { + cooperative_groups::this_grid().sync(); +} + +// Single persistent kernel: gemv+conv -> [barrier] -> delta. Cooperative +// launch. The kB1 instantiation specializes the serving-hot B=1 case: +// batch/split index math collapses to compile-time constants and the fp32 state +// rows for the (single) delta task of each warp are prefetched before the +// barrier so their long-scoreboard latency overlaps the gemv/conv phases (the +// state pool is only written after the barrier, each row by the warp that +// prefetched it). +// +// Padded batch rows: a NEGATIVE state index (vLLM's PAD_SLOT_ID = -1) marks a +// batch row that owns no pool slot -- what a CUDA-graph replay carries in the +// rows between the live request count and the captured batch size. Such a row +// is skipped in every phase that touches a pool (no read and no write of +// conv_state / ssm_state) and its output rows are written as zero, matching +// the fp32 path of gated_delta_rule_decode_pretranspose. The check has to be +// here rather than on the host: reading index VALUES host-side costs a +// device-to-host sync per layer per decode step and is impossible under graph +// capture. Each guard is uniform over the threads that share a batch row +// (warp-uniform in the delta phase, where the warp-wide shuffles below make +// divergence unacceptable), so it costs one predicated branch, not divergence. +// Indices >= P are NOT padding and are not checked -- see the note on +// state_indices in the FFI entry point below. +// +// Aliasing: the op updates both state pools IN PLACE, so the launcher passes +// the same pointer for (conv_state, updated_conv) and for (ssm_state, +// ssm_out). Those four parameters therefore carry no __restrict__ -- the +// pools are read and written through two different parameters, which is +// exactly what restrict promises does not happen, and promising it would let +// the compiler reorder a pool load across a pool store. The remaining +// pointers are genuinely disjoint buffers and keep the qualifier. The read +// path pays for this: loads from the pools can no longer be promoted to +// ld.global.nc. Runtime admission remains opt-in, with FP32 state unchanged. +template +__global__ void gdn_fused_decode_kernel( + const f16* __restrict__ hidden, const f16* __restrict__ w_ba, + const f16* __restrict__ mixed_qkv, const f16* __restrict__ conv_weight, + const f16* __restrict__ conv_bias, const f16* conv_state, + const float* __restrict__ A_log, const f16* __restrict__ dt_bias, + const float* ssm_state, const int* __restrict__ state_indices, float scale, + long state_stride_0, long qkv_stride, long conv_stride_p, + long conv_stride_c, long conv_stride_t, f16* __restrict__ output, + f16* updated_conv, float* ssm_out, float* __restrict__ ba_part, + f16* __restrict__ conv_out, int B) { + // Each CTA owns half a head (8 warps x 8 rows). The exact same gate and + // normalized Q/K were computed eight times. The experimental shared path + // evaluates that unchanged reduction once, without another grid barrier. + __shared__ float head_parameters[SHARED_PARAMETERS ? 2 * D + 3 : 1]; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int nthreads = gridDim.x * blockDim.x; + const int Beff = kB1 ? 1 : B; + + if constexpr (Phase != 2) { + // ---- Phase A1: GEMV partials ---- + // tasks: (split in 0..GEMV_NSPLIT-1) x (b) x (col in 0..N_BA-1). Partials + // are stored split-major per (col, b) so the gate reduction reads them with + // warp-coalesced loads. + long gemv_tasks = (long)GEMV_NSPLIT * Beff * N_BA; + for (long t = tid; t < gemv_tasks; t += nthreads) { + int col = t % N_BA; + long r = t / N_BA; + int b; + int split; + if constexpr (kB1) { + b = 0; + split = (int)r; + } else { + b = r % B; + split = r / B; + } + const f16* hrow = hidden + (long)b * HIDDEN; + float a0 = 0, a1 = 0, a2 = 0, a3 = 0; + int k = split; + for (; k + 3 * GEMV_NSPLIT < HIDDEN; k += 4 * GEMV_NSPLIT) { + a0 += __half2float(hrow[k]) * __half2float(w_ba[(long)k * N_BA + col]); + a1 += __half2float(hrow[k + GEMV_NSPLIT]) * + __half2float(w_ba[(long)(k + GEMV_NSPLIT) * N_BA + col]); + a2 += __half2float(hrow[k + 2 * GEMV_NSPLIT]) * + __half2float(w_ba[(long)(k + 2 * GEMV_NSPLIT) * N_BA + col]); + a3 += __half2float(hrow[k + 3 * GEMV_NSPLIT]) * + __half2float(w_ba[(long)(k + 3 * GEMV_NSPLIT) * N_BA + col]); + } + for (; k < HIDDEN; k += GEMV_NSPLIT) + a0 += __half2float(hrow[k]) * __half2float(w_ba[(long)k * N_BA + col]); + ba_part[((long)col * Beff + b) * GEMV_NSPLIT + split] = + (a0 + a1) + (a2 + a3); + } + + // ---- Phase A2: conv (independent of gemv) ---- + long conv_tasks = (long)Beff * QKV_DIM; + for (long t = tid; t < conv_tasks; t += nthreads) { + int b; + int c; + if constexpr (kB1) { + b = 0; + c = (int)t; + } else { + b = t / QKV_DIM; + c = t % QKV_DIM; + } + int idx = state_indices[b]; + // Padded row: owns no conv-state slot, so neither shift it nor append to + // it. conv_out for this row stays whatever the scratch held -- the delta + // phase skips the same row, so nothing reads it. + if (idx < 0) continue; + // conv_state addressing is stride-parameterized: the pool arrives as a + // logical [P, QKV_DIM, CONV_STATE_LEN] view of either a DS-dense pool + // (strides p,3,1 -> per-thread 3-element rows) or a transposed SD pool + // (strides p,1,QKV_DIM -> fully coalesced across channels, the vLLM + // default). Pure index arithmetic; the update math is identical. + const f16* st = + conv_state + (long)idx * conv_stride_p + (long)c * conv_stride_c; + f16 s0 = st[0], s1 = st[conv_stride_t], s2 = st[2 * conv_stride_t]; + // mixed_qkv rows may be strided (e.g. a view into a wider projection). + f16 xr = mixed_qkv[(long)b * qkv_stride + c]; + const f16* w = conv_weight + (long)c * CONV_WIDTH; + // vLLM's FP16 conv actually emits mul.f16 then cvt.f32.f16, with + // FP32 accumulation. Preserve those load-bearing product roundings; + // widening before the multiply changes this model's recurrent inputs. + float y = conv_bias ? __half2float(conv_bias[c]) : 0.f; + y += __half2float(__hmul(s0, w[0])); + y += __half2float(__hmul(s1, w[1])); + y += __half2float(__hmul(s2, w[2])); + y += __half2float(__hmul(xr, w[3])); + conv_out[(long)b * QKV_DIM + c] = __float2half_rn(siluf(y)); + f16* uc = + updated_conv + (long)idx * conv_stride_p + (long)c * conv_stride_c; + uc[0] = s1; + uc[conv_stride_t] = s2; + uc[2 * conv_stride_t] = xr; + } + } + if constexpr (Phase == 1) return; + + // ---- Pre-barrier prefetch of this warp's first delta task's state rows. + // The state pool is read-only until phase C, and each row is written only + // by the warp that owns (and prefetched) it, so this is race-free. + int gwarp = tid >> 5; + int lane = threadIdx.x & 31; + int nwarps = nthreads >> 5; + long total_rows = (long)Beff * HV * D; + long warps_needed = (total_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP; + + float4 s_pre[ROWS_PER_WARP]; + long pre_row_base = -1; + if (gwarp < warps_needed) { + long first_row = (long)gwarp * ROWS_PER_WARP; + int v0; + int h; + int b; + if constexpr (kB1) { + v0 = (int)(first_row & (D - 1)); + h = (int)(first_row >> D_LOG2); + b = 0; + } else { + v0 = first_row % D; + long tmp = first_row / D; + h = tmp % HV; + b = tmp / HV; + } + int idx = state_indices[b]; + // A padded row has no state row to prefetch. pre_row_base stays -1, which + // no live row_base can equal, so the delta phase never mistakes the + // (unwritten) s_pre registers for this warp's prefetched rows. + if (idx >= 0) { + pre_row_base = (long)idx * state_stride_0 + (long)h * (D * D); + const float4* base_srow = + (const float4*)(ssm_state + pre_row_base + (long)v0 * D); +#pragma unroll + for (int r = 0; r < ROWS_PER_WARP; ++r) + s_pre[r] = base_srow[r * (D / 4) + lane]; + } + } + + // Split-phase launches use stream ordering for this producer/consumer + // boundary and do not reserve the whole cooperative grid. Arithmetic, + // half round trips and ownership of each state row are unchanged. + if constexpr (Phase == 0) grid_barrier(); + + // ---- Phase C: delta (gate reduced inline from ba_part) ---- + for (long w = gwarp; w < warps_needed; w += nwarps) { + long first_row = w * ROWS_PER_WARP; + int v0; + int h; + int b; + if constexpr (kB1) { + v0 = (int)(first_row & (D - 1)); + h = (int)(first_row >> D_LOG2); + b = 0; + } else { + v0 = first_row % D; + long tmp = first_row / D; + h = tmp % HV; + b = tmp / HV; + } + int j = h / HEADS_PER_QK; + const f16* co = conv_out + (long)b * QKV_DIM; + const f16* qb = co + j * D; + const f16* kb = co + H_Q * D + j * D; + int k0 = lane * 4; + int idx = state_indices[b]; + if (idx < 0) { + // Padded row: no state row to read or write; its output rows are zero. + // b (hence idx) is warp-uniform -- first_row is a multiple of + // ROWS_PER_WARP and the whole warp shares w -- so the whole warp takes + // this branch together and the warp-wide shuffles below are never + // reached with a partial mask. + if (lane == 0) { +#pragma unroll + for (int r = 0; r < ROWS_PER_WARP; ++r) + output[((long)b * HV + h) * D + v0 + r] = __float2half_rn(0.0f); + } + continue; + } + long row_base = (long)idx * state_stride_0 + (long)h * (D * D); + float4 s4[ROWS_PER_WARP]; + if (w == gwarp && row_base == pre_row_base) { + // First iteration (the only one for B=1): use the prefetched rows. +#pragma unroll + for (int r = 0; r < ROWS_PER_WARP; ++r) s4[r] = s_pre[r]; + } else { + const float4* base_srow = + (const float4*)(ssm_state + row_base + (long)v0 * D); +#pragma unroll + for (int r = 0; r < ROWS_PER_WARP; ++r) + s4[r] = base_srow[r * (D / 4) + lane]; + } + float qh[4], kh[4], QK, beta, g; + if (!SHARED_PARAMETERS || threadIdx.x < 32) { + // Issue the gate-partial loads early (10 concurrent warp-wide loads) so + // they overlap the qk-norm compute below. + const float* base_b = ba_part + ((long)h * Beff + b) * GEMV_NSPLIT; + const float* base_a = ba_part + ((long)(HV + h) * Beff + b) * GEMV_NSPLIT; + float b0 = base_b[lane + 0]; + float a0v = base_a[lane + 0]; + float b1 = base_b[lane + 32]; + float a1v = base_a[lane + 32]; + float b2 = base_b[lane + 64]; + float a2v = base_a[lane + 64]; + float b3 = base_b[lane + 96]; + float a3v = base_a[lane + 96]; + float b4 = base_b[lane + 128]; + float a4v = base_a[lane + 128]; + float qraw[4], kraw[4]; + float qss = 0.f, kss = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + qraw[i] = __half2float(qb[k0 + i]); + kraw[i] = __half2float(kb[k0 + i]); + qss += qraw[i] * qraw[i]; + kss += kraw[i] * kraw[i]; + } + qss = warp_reduce(qss); + kss = warp_reduce(kss); + qss = __shfl_sync(0xffffffff, qss, 0); + kss = __shfl_sync(0xffffffff, kss, 0); + float qn = rsqrtf(qss + 1e-6f), kn = rsqrtf(kss + 1e-6f); + float QKp = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + qh[i] = qraw[i] * qn; + kh[i] = kraw[i] * kn; + QKp += qh[i] * kh[i]; + } + QKp = warp_reduce(QKp); + QK = __shfl_sync(0xffffffff, QKp, 0); + // gate g,beta reduced from the split-major partials (values now arrived) + float accb = ((b0 + b1) + (b2 + b3)) + b4; + float acca = ((a0v + a1v) + (a2v + a3v)) + a4v; + accb = warp_reduce(accb); + acca = warp_reduce(acca); + accb = __shfl_sync(0xffffffff, accb, 0); + acca = __shfl_sync(0xffffffff, acca, 0); + // The f16 round-trip of the two gate sums is load-bearing, not a + // leftover: the composable path materializes `ba = (hidden @ w_ba)` as a + // f16 tensor and only then widens it for the gates, so the values it + // feeds sigmoid/softplus are f16-rounded. Keeping the fp32 accumulator + // here would make this kernel *more* precise than the operation it + // implements and move the gates off the reference by up to one f16 ulp -- + // amplified by exp() in the decay gate. Track the composable path's `ba` + // dtype, not the accumulator's. + beta = sigmoidf(__half2float(__float2half_rn(accb))); + float xg = __half2float(__float2half_rn(acca)) + __half2float(dt_bias[h]); + g = __expf(-__expf(A_log[h]) * softplusf(xg)); + if constexpr (SHARED_PARAMETERS) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + head_parameters[k0 + i] = qh[i]; + head_parameters[D + k0 + i] = kh[i]; + } + if (lane == 0) { + head_parameters[2 * D] = QK; + head_parameters[2 * D + 1] = beta; + head_parameters[2 * D + 2] = g; + } + } + } + if constexpr (SHARED_PARAMETERS) { + __syncthreads(); +#pragma unroll + for (int i = 0; i < 4; ++i) { + qh[i] = head_parameters[k0 + i]; + kh[i] = head_parameters[D + k0 + i]; + } + QK = head_parameters[2 * D]; + beta = head_parameters[2 * D + 1]; + g = head_parameters[2 * D + 2]; + } +#pragma unroll + for (int r = 0; r < ROWS_PER_WARP; ++r) { + int v = v0 + r; + float s[4] = {s4[r].x, s4[r].y, s4[r].z, s4[r].w}; + float vv = __half2float(co[2 * H_Q * D + h * D + v]); + float kSp = 0.f, qSp = 0.f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + kSp += kh[i] * s[i]; + qSp += qh[i] * s[i]; + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { + kSp += __shfl_down_sync(0xffffffff, kSp, o); + qSp += __shfl_down_sync(0xffffffff, qSp, o); + } + float kS = __shfl_sync(0xffffffff, kSp, 0); + float qS = __shfl_sync(0xffffffff, qSp, 0); + float old_v = g * kS; + float delta = beta * (vv - old_v); + float out_v = scale * (g * qS + delta * QK); + float4* Sorow = (float4*)(ssm_out + row_base + (long)v * D); + float4 o4; + o4.x = g * s[0] + kh[0] * delta; + o4.y = g * s[1] + kh[1] * delta; + o4.z = g * s[2] + kh[2] * delta; + o4.w = g * s[3] + kh[3] * delta; + Sorow[lane] = o4; + if (lane == 0) + output[((long)b * HV + h) * D + v] = __float2half_rn(out_v); + } + // All eight warps have identical loop bounds and padding status. Do not + // let a faster warp's next head overwrite values still being consumed. + if constexpr (SHARED_PARAMETERS) __syncthreads(); + } +} + +} // namespace FI_GDN_IMPL_NAMESPACE diff --git a/flashinfer-sm70/include/flashinfer/hc/sm70/hc_combine_norm.cuh b/flashinfer-sm70/include/flashinfer/hc/sm70/hc_combine_norm.cuh new file mode 100644 index 0000000000..6091aa7188 --- /dev/null +++ b/flashinfer-sm70/include/flashinfer/hc/sm70/hc_combine_norm.cuh @@ -0,0 +1,185 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright (c) 2026 FlashInfer team +// Adapted by 1Cat-vLLM contributors from FlashInfer norm.cuh at +// 6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f. HC injection and materialized +// residual rounding are model semantics absent from upstream fused-add norm. +#pragma once +#include +namespace flashinfer::sm70::hc { +// Same fused-add norm algorithm, but a small fixed-width decode row can +// retain its rounded residual in registers instead of staging/reloading it +// through shared memory. Geometry specialization is local to this operator. +template +__global__ void HCCombineNormRegisterKernel( + const B* __restrict__ input, const T* __restrict__ residual, + const W* __restrict__ weight, const W* __restrict__ injection, + T* __restrict__ combined, T* __restrict__ output, uint32_t groups, + bool shared_weight, float eps) { + constexpr uint32_t THREADS = 32 * WARPS; + constexpr uint32_t ROUNDS = ceil_div(D, VEC_SIZE * THREADS); + const uint32_t row = blockIdx.x, group = blockIdx.y; + const uint32_t tid = threadIdx.x, lane = tid % 32, warp = tid / 32; + const uint32_t offset = (row * groups + group) * D; + const float gate = + 2.f / (1.f + __expf(-float(injection[row * groups + group]) / groups)); + vec_t values[ROUNDS]; + float sum_sq = 0.f; +#pragma unroll + for (uint32_t round = 0; round < ROUNDS; ++round) { + const uint32_t col = (round * THREADS + tid) * VEC_SIZE; + vec_t bv; + vec_t rv; + bv.fill(0.f); + rv.fill(0.f); + if (col < D) { + bv.load(input + row * D + col); + rv.load(residual + offset + col); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + const float x = float(T(fmaf(float(bv[j]), gate, float(rv[j])))); + values[round][j] = x; + rv[j] = T(x); + sum_sq += x * x; + } + if (col < D) rv.store(combined + offset + col); + } +#pragma unroll + for (uint32_t delta = 16; delta > 0; delta /= 2) + sum_sq += math::shfl_xor_sync(sum_sq, delta); + __shared__ float sums[WARPS]; + if (lane == 0) sums[warp] = sum_sq; + __syncthreads(); + if (warp == 0) { + sum_sq = lane < WARPS ? sums[lane] : 0.f; +#pragma unroll + for (uint32_t delta = 16; delta > 0; delta /= 2) + sum_sq += math::shfl_xor_sync(sum_sq, delta); + if (lane == 0) sums[0] = sum_sq; + } + __syncthreads(); + const float inv_rms = math::rsqrt(sums[0] / D + eps); +#pragma unroll + for (uint32_t round = 0; round < ROUNDS; ++round) { + const uint32_t col = (round * THREADS + tid) * VEC_SIZE; + vec_t wv; + vec_t result; + if (col < D) { + wv.load(weight + (shared_weight ? 0 : group * D) + col); +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + const float y = values[round][j] * inv_rms; + result[j] = T(fmaf(y, float(wv[j]), y)); + } + result.store(output + offset + col); + } + } +} + +template +__global__ void HCCombineNormKernel( + const B* __restrict__ input, const T* __restrict__ residual, + const W* __restrict__ weight, const W* __restrict__ injection, + T* __restrict__ combined, T* __restrict__ output, uint32_t groups, + uint32_t d, uint32_t stride_input, uint32_t stride_residual, + uint32_t stride_injection, bool shared_weight, float eps) { + const uint32_t bx = blockIdx.x, group = blockIdx.y; + const float gate = + 2.f / + (1.f + __expf(-float(injection[bx * stride_injection + group]) / groups)); + const uint32_t weight_offset = shared_weight ? 0 : group * d; + const uint32_t tx = threadIdx.x, ty = threadIdx.y; + constexpr uint32_t warp_size = 32; + const uint32_t num_warps = blockDim.y; + const uint32_t thread_id = tx + ty * warp_size; + const uint32_t num_threads = num_warps * warp_size; + const uint32_t rounds = ceil_div(d, VEC_SIZE * num_threads); + extern __shared__ float smem[]; + float* smem_x = smem + ceil_div(num_warps, 4) * 4; + + float sum_sq = 0.f; +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + for (uint32_t i = 0; i < rounds; i++) { + vec_t input_vec; + input_vec.fill(0.f); + vec_t residual_vec; + residual_vec.fill(0.f); + vec_t x_vec; + x_vec.fill(0.f); + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + input_vec.load(input + bx * stride_input + i * num_threads * VEC_SIZE + + thread_id * VEC_SIZE); + residual_vec.load(residual + bx * stride_residual + group * d + + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; j++) { + // Match HC combine -> materialized residual -> Gemma norm rounding. + float x = + float(T(fmaf(float(input_vec[j]), gate, float(residual_vec[j])))); + sum_sq += x * x; + residual_vec[j] = (T)x; + x_vec[j] = x; + } + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + residual_vec.store(combined + bx * stride_residual + group * d + + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + x_vec.store(smem_x + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } + } + + // first, warp reduce sum +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) { + sum_sq += math::shfl_xor_sync(sum_sq, offset); + } + + if (tx == 0) smem[ty] = sum_sq; + __syncthreads(); + // then, cross warp reduce sum using only the first warp + if (ty == 0) { + sum_sq = (tx < num_warps) ? smem[tx] : 0.f; +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) { + sum_sq += math::shfl_xor_sync(sum_sq, offset); + } + if (tx == 0) smem[0] = sum_sq; + } + __syncthreads(); + + float rms_rcp = math::rsqrt(smem[0] / float(d) + eps); + + for (uint32_t i = 0; i < rounds; i++) { + vec_t input_vec; + vec_t weight_vec; + vec_t x_vec; + input_vec.fill(0.f); + weight_vec.fill(0.f); + x_vec.fill(0.f); + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + weight_vec.load(weight + weight_offset + i * num_threads * VEC_SIZE + + thread_id * VEC_SIZE); + x_vec.load(smem_x + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; j++) { + const float y = x_vec[j] * rms_rcp; + input_vec[j] = fmaf(y, float(weight_vec[j]), y); + } + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + input_vec.store(output + bx * stride_residual + group * d + + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +} // namespace flashinfer::sm70::hc diff --git a/flashinfer-sm70/tests/test_batch_bridges.py b/flashinfer-sm70/tests/test_batch_bridges.py new file mode 100644 index 0000000000..18e1cd33a5 --- /dev/null +++ b/flashinfer-sm70/tests/test_batch_bridges.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +from pathlib import Path +from types import SimpleNamespace as NS + +import pytest +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv import ( + capture, + check_exclusive, + load_weights, +) +from benchmarks.kernels.benchmark_sm70_flashinfer_qsa import make_case +from vllm import envs +from vllm.model_executor.layers import sm70_flashinfer_batch as fi +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first + + +@pytest.fixture(scope="module") +def cuda_bridge(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + check_exclusive() + for key in ("VLLM_SM70_FLASHINFER_GDN_LIBRARY", "VLLM_SM70_FLASHINFER_QSA_LIBRARY"): + torch.ops.load_library(os.environ[key]) + zero = torch.zeros(256, device="cuda", dtype=torch.float16) + fi._QSA_ZERO[zero.device] = zero + yield + fi._QSA_ZERO.pop(zero.device) + + +@pytest.mark.parametrize("rows", [4, 8, 9, 15, 16]) +@pytest.mark.parametrize("kv_heads,selected", [(1, 2051), (2, 2051), (1, 15), (1, 65)]) +def test_qsa_bridge_graph_and_materialized_gate( + cuda_bridge, monkeypatch, rows, kv_heads, selected +): + from vllm.models.qwen4_exp.nvidia.ops.qsa import ( + _qsa_output_gate, + qsa_sparse_paged_attention, + ) + + envs.disable_envs_cache() + monkeypatch.setenv("VLLM_SM70_FLASHINFER_BATCH", "0") + torch.manual_seed(37) + q, k, v, indices, table, requests = make_case(rows) + if kv_heads != 1: + q = q.repeat(1, kv_heads, 1) + k = k.repeat(1, 1, kv_heads, 1) + v = v.repeat(1, 1, kv_heads, 1) + indices = indices[:, :selected] + output = torch.empty_like(q) + gate = torch.randn_like(q) + + def call(): + assert fi.try_qsa(q, k, v, indices, table, requests, output) is output + _qsa_output_gate(output, gate) + + graph = capture(call) + for cycle in range(4): + q.normal_() + gate.normal_() + indices[:, :8] = cycle + if cycle == 3: + indices[-1].fill_(-1) + reference = qsa_sparse_paged_attention( + q, k, v, indices, table, requests, output_gate=gate + ) + call() + eager = output.clone() + output.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(output, eager, atol=0, rtol=0) + relative = ( + output.float() - reference.float() + ).norm() / reference.float().norm() + assert torch.isfinite(output).all() and relative < 5e-3 + + +@pytest.mark.parametrize("empty_aot_placeholders", [False, True]) +def test_gdn_bridge_independent_state_and_graph(cuda_bridge, empty_aot_placeholders): + from flash_qla.ops.gated_delta_rule.chunk.sm70 import fused_fwd as qla + from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_update + + torch.manual_seed(17) + model = os.environ.get("SM70_FLASHINFER_TEST_MODEL") + if not model or not Path(model).is_dir(): + pytest.skip("Set SM70_FLASHINFER_TEST_MODEL to the checkpoint directory") + h, hq, hv, wqkv, ba, cw, A, dt = load_weights(Path(model)) + rows, width, pool = 8, wqkv.shape[0], 11 + weight = torch.cat((wqkv, torch.randn(hv * 128, h, device="cuda").half() * 0.01)) + layer = NS( + _sm70_fi_ready=True, + num_v_heads=hv * 4, + num_k_heads=hq * 4, + tp_size=4, + _sm70_fi_ba=ba.t().contiguous(), + _sm70_fi_bias=ba.new_empty(0), + _sm70_fi_op=torch.ops._C_flashinfer_gdn_sm70_h2560_q4_v12.run, + conv1d=NS(weight=cw[:, None, :]), + A_log=A, + dt_bias=dt, + in_proj_qkvz=lambda x: (torch.nn.functional.linear(x, weight), None), + ) + hidden = torch.randn(rows, h, device="cuda", dtype=torch.float16) + conv = torch.randn(pool, width, 3, device="cuda", dtype=torch.float16) * 0.1 + state = torch.randn(pool, hv, 128, 128, device="cuda") * 0.01 + raw_conv = conv if is_conv_state_dim_first() else conv.transpose(-1, -2) + layer.kv_cache = ( + (raw_conv, state) + if empty_aot_placeholders + else (raw_conv.clone(), state.clone()) + ) + placeholder = hidden.new_empty(0) + indices = torch.arange(rows, device="cuda", dtype=torch.int32) + meta = NS( + num_prefills=0, + num_prefill_tokens=0, + num_spec_decodes=0, + num_spec_decode_tokens=0, + num_decodes=rows, + num_decode_tokens=rows, + non_spec_state_indices_tensor=indices, + ) + z = hidden.new_empty(rows, hv, 128) + output = torch.empty_like(z) + ref_out = torch.empty_like(z) + + def call(): + assert fi.try_gdn( + layer, + hidden, + z, + output, + placeholder if empty_aot_placeholders else raw_conv, + placeholder if empty_aot_placeholders else state, + meta, + ) + + graph = capture(call) + ref_conv, ref_state = conv.clone(), state.clone() + for cycle in range(8): + hidden.normal_().mul_(0.5) + indices.copy_(torch.randperm(pool, device="cuda")[:rows]) + if cycle == 7: + indices[-1] = -1 + qkvz = torch.nn.functional.linear(hidden, weight) + raw_qkv = qkvz[:, :width].clone() + ba_ref = hidden @ ba.t() + causal_conv1d_update( + raw_qkv, + ref_conv, + cw, + None, + "silu", + conv_state_indices=indices, + validate_data=False, + ) + qla.gdn_decode_mixed_qkv_global_state_sm70( + raw_qkv, + ba_ref[:, hv:].contiguous(), + ba_ref[:, :hv].contiguous(), + A, + dt, + ref_state, + indices, + ref_out, + ) + old_c, old_s = conv.clone(), state.clone() + call() + expected = [t.clone() for t in (z, output, conv, state)] + conv.copy_(old_c) + state.copy_(old_s) + output.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + for actual, eager in zip((z, output, conv, state), expected): + torch.testing.assert_close(actual, eager, atol=0, rtol=0) + torch.testing.assert_close(z, qkvz[:, width:].reshape_as(z), atol=0, rtol=0) + torch.testing.assert_close(conv, ref_conv, atol=0, rtol=0) + assert (state - ref_state).norm() / ref_state.norm() < 5e-3 + live = indices >= 0 + relative = (output[live].float() - ref_out[live].float()).norm() + assert relative / ref_out[live].float().norm() < 5e-3 diff --git a/flashinfer-sm70/tests/test_batch_routing.py b/flashinfer-sm70/tests/test_batch_routing.py new file mode 100644 index 0000000000..741dfbb5a7 --- /dev/null +++ b/flashinfer-sm70/tests/test_batch_routing.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace as NS + +import pytest +import torch + +from vllm import envs +from vllm.model_executor.layers import sm70_flashinfer_batch as fi + + +def metadata(rows=8, **changes): + fields = dict( + num_prefills=0, + num_prefill_tokens=0, + num_spec_decodes=0, + num_spec_decode_tokens=0, + num_decodes=rows, + num_decode_tokens=rows, + non_spec_state_indices_tensor=torch.arange(rows, dtype=torch.int32), + ) + fields.update(changes) + return NS(**fields) + + +def test_opt_in_only(monkeypatch): + envs.disable_envs_cache() + monkeypatch.delenv("VLLM_SM70_FLASHINFER_BATCH", raising=False) + assert not envs.VLLM_SM70_FLASHINFER_BATCH + + +@pytest.mark.parametrize("rows", [2, 4, 8, 16, 32, 64]) +def test_decode_shape_is_not_a_server_configuration(rows): + assert fi.uniform_decode(metadata(rows), rows) + + +@pytest.mark.parametrize( + "changes", + [ + dict(num_prefills=1), + dict(num_prefill_tokens=1), + dict(num_spec_decodes=1), + dict(num_spec_decode_tokens=1), + dict(num_decode_tokens=9), + dict(non_spec_state_indices_tensor=None), + ], +) +def test_prefill_spec_and_invalid_metadata_fall_back(changes): + assert not fi.uniform_decode(metadata(**changes), 8) + + +def test_padding_uses_scheduler_owned_slots_without_host_values(): + meta = metadata(5, non_spec_state_indices_tensor=torch.full((8,), -1)) + assert fi.uniform_decode(meta, 8) + assert not fi.uniform_decode(meta, 16) + + +def test_m1_and_unknown_keep_original_paths(): + assert not fi.uniform_decode(metadata(1), 1) + assert not fi.uniform_decode(None, 8) + assert not fi.uniform_decode(metadata(65), 65) + + +def test_unprepared_gdn_does_not_touch_projection_or_state(): + assert not fi.try_gdn(NS(), torch.empty(8, 2560), None, None, None, None, None) + + +def test_unbound_aot_placeholders_fall_back_before_transpose(): + layer = NS(_sm70_fi_ready=True) + empty = torch.empty(0) + assert not fi.try_gdn( + layer, torch.empty(8, 2560), None, None, empty, empty, metadata() + ) + + +def test_unprepared_qsa_does_not_allocate_workspace(): + q = torch.empty(8, 6, 256, dtype=torch.float16) + assert fi.try_qsa(q, None, None, None, None, None, None) is None + + +@pytest.mark.parametrize("rows", [4, 8, 9, 15, 16]) +@pytest.mark.parametrize("compatible", [False, True]) +@pytest.mark.parametrize("kv_heads,selected", [(1, 2051), (2, 2051), (1, 15), (1, 65)]) +def test_qsa_compatibility_owns_partition_and_call_local_scratch( + monkeypatch, rows, compatible, kv_heads, selected +): + calls = [] + native = NS(run=lambda *args: calls.append(args)) + monkeypatch.setattr(torch.ops, "_C_flashinfer_qsa_sm70", native) + monkeypatch.setattr( + torch.ops, "_C_flashinfer_qsa_sm70_compat", native if compatible else NS() + ) + q = torch.empty(rows, 6 * kv_heads, 256, dtype=torch.float16) + k = torch.empty(rows, 64, kv_heads, 256, dtype=q.dtype) + indices = torch.empty(rows, selected, dtype=torch.int32) + table = torch.empty(rows, 1, dtype=torch.int32) + requests = torch.arange(rows, dtype=torch.int32) + out = torch.empty_like(q) + zero = torch.zeros(256, dtype=q.dtype) + monkeypatch.setitem(fi._QSA_ZERO, q.device, zero) + for _ in range(2): + assert fi.try_qsa(q, k, k, indices, table, requests, out) is out + if compatible: + programs = rows * kv_heads + target = 64 if programs <= 8 else 32 if programs < 32 else 8 + splits = min(target, 1 << (((selected + 15) // 16).bit_length() - 1)) + else: + splits = 16 if rows >= 16 else 32 + assert len(calls) == 2 + assert calls[0][-1] == splits + assert calls[0][8] is zero and calls[0][11] is out + assert calls[0][9].shape == (rows, splits, 6 * kv_heads, 256) + for index in (6, 7, 9, 10): + assert calls[0][index].data_ptr() != calls[1][index].data_ptr() + + +def test_unprepared_mqa_does_not_allocate_or_read_metadata(): + q = torch.empty(8, 4, 128, dtype=torch.float16) + assert not fi.try_mqa(q, None, None, None, None, None, 4, 1.0, None, None) + + +@pytest.mark.parametrize("rows,workers", [(4, 160), (8, 160), (16, 320)]) +def test_mqa_dispatch_preserves_int64_positions_and_caller_outputs( + monkeypatch, rows, workers +): + # CPU-only dispatch test: substitute the op, not a CUDA implementation. + device = torch.device("cpu") + monkeypatch.setitem(fi._MQA_SMS, device, 80) + seen = [] + monkeypatch.setattr( + torch.ops._C_flashinfer_mqa_sm70, + "run", + lambda *args: seen.append(args), + raising=False, + ) + q = torch.empty(rows, 4, 128, dtype=torch.float16) + k = torch.empty(rows, 196, 1, 128, dtype=torch.float16) + table = torch.empty(rows, 1, dtype=torch.int32) + requests = torch.arange(rows, dtype=torch.int32) + positions = torch.arange(rows, dtype=torch.int64) + lengths = torch.full_like(requests, 196) + out = torch.empty(rows, 196) + visible = torch.empty_like(requests) + assert fi.try_mqa( + q, k, table, requests, positions, lengths, 4, 128**0.5, out, visible + ) + assert len(seen) == 1 + assert seen[0][4] is positions + assert seen[0][6] is out and seen[0][7] is visible + assert seen[0][-1] == workers + assert seen[0][8].shape == (workers + 1, 2) + # Other dtype/layout/components fall back locally, without attempting CUDA. + assert not fi.try_mqa( + q.float(), k, table, requests, positions, lengths, 4, 1.0, out, visible + ) + assert not fi.try_mqa( + q, k, table, requests.long(), positions, lengths, 4, 1.0, out, visible + ) + assert len(seen) == 1 + + +def test_cpu_prepare_is_a_noop(monkeypatch): + monkeypatch.setattr( + torch.ops, "load_library", lambda *_: pytest.fail("unexpected load") + ) + fi.prepare(torch.nn.Module(), torch.device("cpu")) + + +@pytest.mark.parametrize("existing", ["first", "second"]) +def test_preloaded_component_prevents_duplicate_fragment_registration( + monkeypatch, existing +): + from types import SimpleNamespace + + fake_ops = SimpleNamespace(first=SimpleNamespace(), second=SimpleNamespace()) + getattr(fake_ops, existing).run = lambda: None + monkeypatch.setattr(torch, "ops", fake_ops) + monkeypatch.setattr( + fi.importlib.util, "find_spec", lambda *_: pytest.fail("duplicate lookup") + ) + fi.load_native_fragment("unused", ("first", "second")) + + +@pytest.mark.parametrize("origin", [None, "/package/vllm/native.abi3.so"]) +def test_native_fragment_missing_or_wheel_resolved(monkeypatch, origin): + from types import SimpleNamespace + + loaded = [] + monkeypatch.setattr( + torch, + "ops", + SimpleNamespace(first=SimpleNamespace(), load_library=loaded.append), + ) + monkeypatch.setattr( + fi.importlib.util, + "find_spec", + lambda _: SimpleNamespace(origin=origin) if origin else None, + ) + fi.load_native_fragment("unused", ("first",)) + assert loaded == ([origin] if origin else []) + + +@pytest.mark.parametrize("prepared", [False, True]) +def test_unsupported_gdn_reload_cannot_keep_stale_derived_weights( + monkeypatch, prepared +): + from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn + + class UnsupportedGDN(torch.nn.Module): + def __init__(self): + super().__init__() + self.num_k_heads, self.num_v_heads, self.tp_size = 16, 48, 4 + self.in_proj_ba = NS(weight=None) + self.in_proj_qkvz = NS(weight=None) + self._sm70_fi_ready = prepared + + for key in ("VLLM_SM70_FLASHINFER_GDN_LIBRARY", "VLLM_SM70_FLASHINFER_QSA_LIBRARY"): + monkeypatch.delenv(key, raising=False) + monkeypatch.setattr( + qwen_gdn_linear_attn, "QwenGatedDeltaNetAttention", UnsupportedGDN + ) + monkeypatch.setattr(fi, "current_platform", NS(is_device_capability=lambda _: True)) + monkeypatch.setattr( + fi, "get_current_vllm_config", lambda: NS(speculative_config=None) + ) + monkeypatch.setattr(fi, "load_native_fragment", lambda *args: None) + layer = UnsupportedGDN() + if prepared: + with pytest.raises(RuntimeError, match="rebuild CUDA graphs"): + fi.prepare(layer, torch.device("cuda")) + else: + fi.prepare(layer, torch.device("cuda")) + + +def test_weight_reload_preserves_captured_pointer(): + module = torch.nn.Module() + weight = torch.randn(5, 3).half() + fi.copy_derived_buffer(module, "packed", weight.t()) + pointer = module.packed.data_ptr() + weight.fill_(2) + fi.copy_derived_buffer(module, "packed", weight.t()) + assert module.packed.data_ptr() == pointer + assert not module.state_dict() + torch.testing.assert_close(module.packed, weight.t()) + with pytest.raises(RuntimeError, match="rebuild CUDA graphs"): + fi.copy_derived_buffer(module, "packed", torch.empty(4, 5)) diff --git a/flashinfer-sm70/tests/test_compiled_gdn_boundary.py b/flashinfer-sm70/tests/test_compiled_gdn_boundary.py new file mode 100644 index 0000000000..5c0c655954 --- /dev/null +++ b/flashinfer-sm70/tests/test_compiled_gdn_boundary.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A large prefill trace must not erase the small-decode runtime dispatch.""" + +import ast +import inspect +import textwrap +from contextlib import nullcontext +from types import SimpleNamespace as NS + +import pytest +import torch + +from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gdn + + +@pytest.mark.parametrize("legacy_shape_guard", [False, True]) +def test_boundary_exports_the_whole_prefill_decode_range(legacy_shape_guard): + # Export the actual first branch predicate, not a duplicate of its logic. + # The old `1 < num_tokens <= 64` condition constrains this graph to rows + # >=65 and fails export's requested 1..2048 contract after a prefill trace. + source = ast.parse( + textwrap.dedent(inspect.getsource(gdn.QwenGatedDeltaNetAttention.forward_cuda)) + ) + branch = next(node for node in source.body[0].body if isinstance(node, ast.If)) + probe = ast.parse( + "def forward(self, x):\n" + " num_tokens = x.size(0)\n" + " if True:\n" + " return x + 1\n" + " return x - 1\n" + ) + probe.body[0].body[1].test = branch.test + if legacy_shape_guard: + probe.body[0].body[1].test = ast.BoolOp( + op=ast.And(), + values=[branch.test, ast.parse("1 < num_tokens <= 64", mode="eval").body], + ) + scope = { + "_sm70_qwen_gdn_input_core_boundary_enabled": lambda: False, + "use_sm70_decode_graph_semantics": lambda: True, + } + exec(compile(ast.fix_missing_locations(probe), __file__, "exec"), scope) + module = type("BoundaryProbe", (torch.nn.Module,), {"forward": scope["forward"]})() + module._sm70_fi_ready = True + expected_error = ( + pytest.raises(torch._dynamo.exc.UserError, match="Constraints violated") + if legacy_shape_guard + else nullcontext() + ) + with expected_error: + exported = torch.export.export( + module, + (torch.ones(128, 2),), + dynamic_shapes={"x": {0: torch.export.Dim("rows", min=1, max=2048)}}, + ).module() + if legacy_shape_guard: + return + for rows in (1, 4, 16, 128, 2048): + torch.testing.assert_close( + exported(torch.ones(rows, 2)), torch.full((rows, 2), 2.0) + ) + + +@pytest.mark.parametrize("rows", [1, 4, 65]) +def test_unsupported_runtime_keeps_existing_fp16_projection(monkeypatch, rows): + layer = NS( + _sm70_fi_ready=True, + sm70_qwen38_fp16_fused_input=True, + in_proj_qkvz=NS(weight=torch.empty(0)), + in_proj_ba=NS(weight=torch.empty(0)), + ) + ctx = NS(no_compile_layers={"layer": layer}, attn_metadata=None) + monkeypatch.setattr(gdn, "get_forward_context", lambda: ctx) + monkeypatch.setattr(gdn, "_resolve_layer_name", lambda name: name) + monkeypatch.setattr( + gdn, "_sm70_dump_gdn_projection_tensor", lambda _, __, tensor: tensor + ) + monkeypatch.setattr(gdn, "_sm70_gdn_projection_dump_requested", lambda _: False) + monkeypatch.setattr(gdn, "_sm70_gdn_qpn8_ba_split_eligible", lambda *_: False) + monkeypatch.setattr(gdn, "use_sm70_decode_graph_semantics", lambda: True) + qkv = torch.randn(rows, 5) + z = torch.randn(rows, 6) + b, a = torch.randn(rows, 3), torch.randn(rows, 3) + calls = [] + + def project(*args): + calls.append("original_projection") + return qkv, z, b, a + + def core(self, **kwargs): + assert kwargs["mixed_qkv"] is qkv + assert kwargs["b"] is b and kwargs["a"] is a + calls.append("original_core") + kwargs["core_attn_out"].fill_(2) + + monkeypatch.setattr( + torch.ops.vllm, "qwen38_sm70_fp16_gdn_input", project, raising=False + ) + monkeypatch.setattr(gdn, "_qwen_gdn_run_recurrent_core", core) + z_out = torch.empty(rows, 3, 2) + core_out = torch.empty_like(z_out) + result = gdn.qwen_gdn_input_projection_core( + torch.randn(rows, 7), z_out, core_out, None, None, "layer" + ) + assert calls == ["original_projection", "original_core"] + assert result[0] is z_out and result[1] is core_out + torch.testing.assert_close(z_out.flatten(1), z) + torch.testing.assert_close(core_out, torch.full_like(core_out, 2)) diff --git a/flashinfer-sm70/tests/test_layer_fusion.py b/flashinfer-sm70/tests/test_layer_fusion.py new file mode 100644 index 0000000000..3b7134dc1f --- /dev/null +++ b/flashinfer-sm70/tests/test_layer_fusion.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_gdn_conv import ( + capture, + check_exclusive, +) +from benchmarks.kernels.flashinfer_sm70_gdn_conv import FusedGDN, build +from benchmarks.kernels.flashinfer_sm70_hc_norm import HCNorm +from benchmarks.kernels.flashinfer_sm70_hc_norm import build as build_hc + + +def test_conv_product_rounding_is_not_fp32_multiply(): + x = torch.tensor([1.0009765625], dtype=torch.float16) + assert (x * x).float().item() != (x.float() * x.float()).item() + + +@pytest.mark.parametrize("rows", [0, -1, 65]) +def test_invalid_rows_fail_before_gpu_allocation(rows): + with pytest.raises(ValueError): + FusedGDN(rows, device="cpu") + + +@pytest.mark.parametrize("hq,hv", [(0, 12), (4, 0), (4, 7), (-1, 12)]) +def test_invalid_heads_fail_before_gpu_allocation(hq, hv): + with pytest.raises(ValueError): + FusedGDN(4, hq, hv, device="cpu") + + +@pytest.fixture(scope="module") +def cuda_build(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + check_exclusive() + torch.manual_seed(7) + build() + + +@pytest.fixture(scope="module") +def cuda_hc_build(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + check_exclusive() + build_hc() + + +def oracle(x, weights, qkv, cw, conv, A, dt, state, indices): + ba = (x @ weights).half().float() + output = torch.zeros(x.shape[0], 12, 128, device=x.device, dtype=torch.float32) + for row in range(x.shape[0]): + slot = int(indices[row]) + if slot < 0: + continue + # Explicit FP16 product boundary observed in production conv PTX. + y = (conv[slot, :, 0] * cw[:, 0]).float() + y += (conv[slot, :, 1] * cw[:, 1]).float() + y += (conv[slot, :, 2] * cw[:, 2]).float() + y += (qkv[row] * cw[:, 3]).float() + mixed = (y / (1 + (-y).exp())).half().float() + conv[slot, :, :2] = conv[slot, :, 1:].clone() + conv[slot, :, 2] = qkv[row] + q = mixed[:512].reshape(4, 128).repeat_interleave(3, 0) + k = mixed[512:1024].reshape(4, 128).repeat_interleave(3, 0) + v = mixed[1024:].reshape(12, 128) + q *= (q.square().sum(-1, keepdim=True) + 1e-6).rsqrt() + k *= (k.square().sum(-1, keepdim=True) + 1e-6).rsqrt() + decay = ( + -A.exp() * torch.nn.functional.softplus(ba[row, 12:] + dt.float()) + ).exp() + beta = ba[row, :12].sigmoid() + old = state[slot] * decay[:, None, None] + delta = (v - (old * k[:, None, :]).sum(-1)) * beta[:, None] + state[slot] = old + delta[:, :, None] * k[:, None, :] + output[row] = (state[slot] * q[:, None, :]).sum(-1) / (128**0.5) + return output + + +@pytest.mark.parametrize( + "rows,sd_layout", [(1, True), (4, False), (8, True), (16, True)] +) +def test_gdn_graph_dynamic_slots_and_history(cuda_build, rows, sd_layout): + pool, width = rows + 2, 2560 + x = torch.randn(rows, 2560, device="cuda", dtype=torch.float16) + weights = torch.randn(2560, 24, device="cuda", dtype=torch.float16) * 0.01 + qkv_storage = torch.randn(rows, 4096, device="cuda", dtype=torch.float16) + qkv = qkv_storage[:, :width] + cw = torch.randn(width, 4, device="cuda", dtype=torch.float16) * 0.2 + bias = torch.empty(0, device="cuda", dtype=torch.float16) + conv = torch.randn(pool, width, 3, device="cuda", dtype=torch.float16) * 0.2 + if sd_layout: + conv = conv.transpose(1, 2).contiguous().transpose(1, 2) + state = torch.randn(pool, 12, 128, 128, device="cuda", dtype=torch.float32) * 0.02 + A = torch.randn(12, device="cuda", dtype=torch.float32) + dt = torch.randn(12, device="cuda", dtype=torch.float16) + indices = torch.arange(rows, device="cuda", dtype=torch.int32) + candidate = FusedGDN(rows) + call = lambda: candidate(x, weights, qkv, cw, bias, conv, A, dt, state, indices) + graph = capture(call) + expected_c, expected_s = conv.clone(), state.clone() + for cycle in range(8): + x.normal_() + qkv.normal_() + indices.copy_(torch.randperm(pool, device="cuda")[:rows]) + if cycle % 3 == 2: + indices[-1] = -1 + c0, s0 = conv.clone(), state.clone() + # The oracle retains its own history: copying candidate state into it + # each cycle would hide accumulated recurrent drift. + expected_o = oracle(x, weights, qkv, cw, expected_c, A, dt, expected_s, indices) + eager = call().clone() + eager_state, eager_conv = state.clone(), conv.clone() + conv.copy_(c0) + state.copy_(s0) + candidate.partial.fill_(float("nan")) + candidate.output.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(candidate.output, eager, atol=0, rtol=0) + torch.testing.assert_close(state, eager_state, atol=0, rtol=0) + torch.testing.assert_close(conv, eager_conv, atol=0, rtol=0) + torch.testing.assert_close(conv, expected_c, atol=0, rtol=0) + torch.testing.assert_close(state, expected_s, atol=1e-4, rtol=2e-3) + torch.testing.assert_close( + candidate.output.float(), expected_o, atol=1e-4, rtol=2e-3 + ) + + +@pytest.mark.parametrize( + "rows,groups,width,r_dtype,b_dtype,w_dtype,shared,warps", + [ + (4, 4, 2560, torch.float16, torch.float16, torch.float16, False, 2), + (8, 4, 320, torch.float32, torch.float16, torch.float16, True, 1), + (16, 2, 520, torch.float16, torch.float32, torch.float32, False, 4), + (1, 3, 4096, torch.float32, torch.float32, torch.float32, True, 8), + ], +) +def test_hc_graph_rounding_and_shared_weights( + cuda_hc_build, rows, groups, width, r_dtype, b_dtype, w_dtype, shared, warps +): + check_hc_graph(rows, groups, width, r_dtype, b_dtype, w_dtype, shared, warps) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("warps", [4, 8]) +def test_hc_register_graph(cuda_hc_build, dtype, warps): + check_hc_graph( + 16, 4, 2560, dtype, torch.float16, torch.float16, True, warps, registers=True + ) + + +def check_hc_graph( + rows, groups, width, r_dtype, b_dtype, w_dtype, shared, warps, registers=False +): + torch.manual_seed(19) + r = torch.randn(rows, groups * width, device="cuda", dtype=r_dtype) + b = torch.randn(rows, width, device="cuda", dtype=b_dtype) + inj = torch.randn(rows, groups, device="cuda", dtype=w_dtype) + w = torch.randn(width if shared else groups * width, device="cuda", dtype=w_dtype) + candidate = HCNorm(r, warps, registers=registers) + call = lambda: candidate(r, b, inj, w) + graph = capture(call) + for scale in (0.25, 1.0, 3.0): + r.normal_().mul_(scale) + b.normal_().mul_(scale) + inj.normal_() + # Independent wider-precision oracle, with the same explicit + # materialized residual boundary; never normalize the unrounded sum. + gate = 2 * (inj.double() / groups).sigmoid() + combined = ( + r.double().reshape(rows, groups, width) + + b.double()[:, None, :] * gate[:, :, None] + ).to(r_dtype) + y = combined.double() + y *= (y.square().mean(-1, keepdim=True) + 1e-6).rsqrt() + y *= 1 + w.double().reshape(1, 1 if shared else groups, width) + eager = [t.clone() for t in call()] + candidate.combined.fill_(float("nan")) + candidate.normalized.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + for actual, ref in zip((candidate.combined, candidate.normalized), eager): + torch.testing.assert_close(actual, ref, atol=0, rtol=0) + atol = 2e-3 if r_dtype == torch.float16 else 2e-5 + for actual, ref in zip(eager, (combined, y.to(r_dtype))): + torch.testing.assert_close( + actual.reshape(rows, groups, width), ref, atol=atol, rtol=2e-3 + ) diff --git a/flashinfer-sm70/tests/test_paired_stats.py b/flashinfer-sm70/tests/test_paired_stats.py new file mode 100644 index 0000000000..ad614d7775 --- /dev/null +++ b/flashinfer-sm70/tests/test_paired_stats.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + +from benchmarks.kernels.sm70_paired_stats import paired_latency_interval + + +def test_identical_ratio_has_exact_interval(): + result = paired_latency_interval([10, 20, 30, 40, 50], [5, 10, 15, 20, 25]) + assert result["reduction_pct_ci95"] == pytest.approx([50.0, 50.0]) + assert result["positive_lower_bound"] + + +def test_noise_is_not_a_pass(): + assert not paired_latency_interval([10] * 5, [8, 12, 8, 12, 10])[ + "positive_lower_bound" + ] + assert not paired_latency_interval([10] * 5, [10] * 5)["positive_lower_bound"] + + +@pytest.mark.parametrize( + "samples", [[0] * 5, [-1] * 5, [float("nan")] * 5, [float("inf")] * 5, [1] * 4] +) +def test_invalid_samples_fail_closed(samples): + with pytest.raises(ValueError): + paired_latency_interval([10] * 5, samples) diff --git a/flashinfer-sm70/tests/test_qsa_decode.py b/flashinfer-sm70/tests/test_qsa_decode.py new file mode 100644 index 0000000000..6dbbf82948 --- /dev/null +++ b/flashinfer-sm70/tests/test_qsa_decode.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Isolated upstream FlashInfer SM70 adapter tests, no vLLM binary required.""" + +import pytest +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_qsa import ( + capture, + check_exclusive, + make_case, + oracle, +) +from benchmarks.kernels.flashinfer_sm70_qsa import FlashInferQSA, build + + +@pytest.mark.parametrize("width,splits", [(0, 1), (2, 0), (2, 65), (-1, 4)]) +def test_invalid_plan(width, splits): + with pytest.raises(ValueError): + FlashInferQSA(torch.empty((1, 6, 256), dtype=torch.float16), width, splits) + + +def test_oracle_keeps_repeats_and_masks_empty_rows(): + q = torch.zeros(2, 1, 256) + k = torch.zeros(1, 4, 1, 256) + v = torch.zeros_like(k) + v[0, 1] = 3 + indices = torch.tensor([[0, 1, 1, -1], [0, 1, 2, 3]], dtype=torch.int32) + result = oracle(q, k, v, indices, torch.tensor([[0]]), torch.tensor([0, -1])) + torch.testing.assert_close(result[0], torch.full_like(result[0], 2)) + assert result[1].count_nonzero() == 0 + + +@pytest.fixture(scope="module") +def cuda_build(): + if not torch.cuda.is_available(): + pytest.skip("SM70 GPU required") + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("Prototype deliberately restricted to SM70") + check_exclusive() + torch.manual_seed(42) + build() + + +@pytest.mark.parametrize( + "rows,selected,page,kv_heads,group,splits", + [ + (1, 1, 16, 1, 1, 1), + (1, 5, 16, 1, 2, 8), + (2, 31, 64, 2, 4, 4), + (4, 33, 16, 1, 6, 32), + (8, 257, 784, 1, 6, 16), + (16, 2051, 784, 1, 6, 32), + (3, 2051, 64, 2, 6, 64), + (2, 37, 16, 2, 8, 4), + ], +) +def test_sparse_eager_and_graph( + cuda_build, rows, selected, page, kv_heads, group, splits +): + case = list(make_case(rows, selected, page, kv_heads, group)) + q, k, v, indices, table, requests = case + # Strided queries/cache and table, while maintaining vector alignment. + q_storage = torch.empty(rows, q.shape[1] * 2, 256, device="cuda", dtype=q.dtype) + case[0] = q_storage[:, ::2].copy_(q) + case[1] = k.transpose(1, 2).contiguous().transpose(1, 2) + case[2] = v.transpose(1, 2).contiguous().transpose(1, 2) + q, k, v = case[:3] + candidate = FlashInferQSA(q, selected, splits) + call = lambda: candidate(*case) + graph = capture(call, 1) + for cycle in range(5): + q.normal_().mul_((0.25, 1.0, 3.0, 1.0, 1.0)[cycle]) + requests.copy_(torch.randperm(rows, device="cuda")) + table.copy_(torch.randperm(k.shape[0], device="cuda").reshape_as(table)) + indices.random_(0, 8192) + if selected > 1: + indices[:, 1] = indices[:, 0] # preserve duplicate weighting + if cycle == 1: + indices[:, ::3] = -1 + elif cycle == 2: + indices[:, ::3] = table.shape[1] * page + 5 + table[:, 0] = k.shape[0] # invalid physical page + elif cycle == 3: + requests[0] = rows + 1 + elif cycle == 4: + indices.fill_(-1) + k.fill_(float("nan")) + v.fill_(float("nan")) # invalid slots must load zero, not page 0 + expected = oracle(*case) + eager = call().clone() + for workspace in (candidate.partial, candidate.lse, candidate.output): + workspace.fill_(float("nan")) + candidate.offsets.fill_(-7) + candidate.metadata.fill_(-9) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(candidate.output, eager, atol=0, rtol=0) + torch.testing.assert_close( + candidate.output.float(), expected, atol=2e-3, rtol=1e-2 + ) + + +def test_reject_unaligned_inputs(cuda_build): + case = list(make_case(1, selected=5)) + q = case[0] + storage = torch.empty(q.numel() + 1, dtype=q.dtype, device=q.device) + case[0] = storage[1:].view_as(q).copy_(q) + candidate = FlashInferQSA(case[0], 5, 1) + with pytest.raises(RuntimeError, match="aligned"): + candidate(*case) diff --git a/setup.py b/setup.py index 563c148cd9..dae5b24e22 100644 --- a/setup.py +++ b/setup.py @@ -1243,6 +1243,8 @@ def _read_requirements(filename: str) -> list[str]: if _is_cuda(): if _cuda_arch_contains(7, 0): ext_modules.append(CMakeExtension(name="vllm._sm70_sampler_C")) + ext_modules.append(CMakeExtension(name="vllm._sm70_flashinfer_C")) + ext_modules.append(CMakeExtension(name="vllm._sm70_flashinfer_gdn_C")) build_sm70_fa2 = _cuda_arch_contains(7, 0) and not _cuda_arch_at_least(8, 0) if _cuda_arch_at_least(8, 0) or build_sm70_fa2: ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) @@ -1326,6 +1328,7 @@ def _read_requirements(filename: str) -> list[str]: "third_party/deep_gemm/include/**/*.cuh", "third_party/deep_gemm/include/**/*.h", "third_party/deep_gemm/include/**/*.hpp", + "third_party/flashinfer_sm70/NOTICE.txt", ], } diff --git a/tests/benchmarks/test_sm70_batch_tool_quality.py b/tests/benchmarks/test_sm70_batch_tool_quality.py new file mode 100644 index 0000000000..d76abcba94 --- /dev/null +++ b/tests/benchmarks/test_sm70_batch_tool_quality.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading + +import pytest + +from benchmarks.benchmark_sm70_batch_tool_quality import run_cases, score_case +from benchmarks.benchmark_sm70_tool_protocol import _bfcl_argument_matches + +pytestmark = pytest.mark.skip_global_cleanup + + +def schema_case(index): + return { + "id": str(index), + "suite": "json_schema", + "schema": { + "type": "object", + "properties": {"value": {"type": "integer"}}, + "required": ["value"], + }, + "request": {"messages": [{"role": "user", "content": str(index)}]}, + } + + +def good_response(): + return {"ok": True, "finish_reason": "stop", "content": '{"value": 3}'} + + +@pytest.mark.parametrize("concurrency", (1, 4, 8, 16)) +def test_requests_really_overlap_and_keep_dataset_order_and_seeds(concurrency): + barrier = threading.Barrier(concurrency, timeout=10) + + def request(url, payload): + barrier.wait() + return good_response() + + cases = [schema_case(i) for i in range(concurrency)] + result = run_cases(cases, "unused", {"seed": 17}, concurrency, request) + assert result["peak_inflight_client_requests"] == concurrency + assert [r["payload"]["seed"] for r in result["cases"]] == list( + range(17, 17 + concurrency) + ) + assert [r["id"] for r in result["cases"]] == [str(i) for i in range(concurrency)] + assert result["suites"]["json_schema"] == { + "correct": concurrency, + "total": concurrency, + } + assert all("seed" not in c["request"] for c in cases) + + +def test_transport_failure_is_retained_and_not_retried(): + calls = [] + + def request(url, payload): + calls.append(1) + raise TimeoutError("test") + + result = run_cases([schema_case(0)], "unused", {"seed": 0}, 1, request) + assert len(calls) == 1 + assert result["suites"]["json_schema"] == {"correct": 0, "total": 1} + assert "TimeoutError" in result["cases"][0]["response"]["error"] + + +@pytest.mark.parametrize("finish_reason", (None, "length", "abort")) +def test_truncated_or_unfinished_valid_json_does_not_pass(finish_reason): + response = {**good_response(), "finish_reason": finish_reason} + assert score_case(schema_case(0), response) + + +@pytest.mark.parametrize("content", ('{"value": "bad"}', "{}", "not json")) +def test_structurally_invalid_json_does_not_pass(content): + assert score_case(schema_case(0), {**good_response(), "content": content}) + + +def test_bfcl_uses_tool_name_and_arguments_not_just_valid_json(): + case = { + "suite": "bfcl/simple_python", + "entry": { + "function": [ + { + "name": "sum", + "parameters": { + "type": "object", + "required": ["a"], + "properties": {"a": {"type": "integer"}}, + }, + } + ], + }, + "ground_truth": [{"sum": {"a": [3]}}], + "irrelevance": False, + } + function = {"name": "sum", "arguments": '{"a": 3}'} + response = { + "ok": True, + "finish_reason": "tool_calls", + "tool_calls": [{"function": function}], + } + assert not score_case(case, response) + response["finish_reason"] = "stop" + assert score_case(case, response) + response["finish_reason"] = "tool_calls" + function["name"] = '{"name": "sum"}' + assert score_case(case, response) + + +@pytest.mark.parametrize( + "value,expected,valid", + [ + ({"min": 3, "max": 4}, [{"min": [3], "max": [4]}], True), + ({"min": 3}, [{"min": [2, 3], "max": [4, ""]}], True), + ({"name": "New-York"}, [{"name": ["new york"]}], True), + ({"values": [1, 2]}, [{"values": [[1, 2]]}], True), + ({"values": [2, 1]}, [{"values": [[1, 2]]}], False), + ({"min": [3]}, [{"min": [3]}], False), + ({"min": 5}, [{"min": [3, 4]}], False), + ({"min": 3}, [{"min": [3], "max": [4]}], False), + ({"min": 3, "extra": 4}, [{"min": [3]}], False), + ({"min": 3}, ["", {"min": [3]}], True), + ({"min": 3}, [""], False), + ], +) +def test_bfcl_dictionary_alternatives_not_literal_values(value, expected, valid): + assert _bfcl_argument_matches(value, expected, {"type": "dict"}) == valid + + +def test_bfcl_list_of_dicts_preserves_order_count_and_literal_array_semantics(): + schema = {"type": "array", "items": {"type": "dict"}} + allowed = [[{"a": [1]}, {"a": [2]}]] + assert _bfcl_argument_matches([{"a": 1}, {"a": 2}], allowed, schema) + assert not _bfcl_argument_matches([{"a": 2}, {"a": 1}], allowed, schema) + assert not _bfcl_argument_matches([{"a": 1}], allowed, schema) + assert not _bfcl_argument_matches([{"a": 1}, 2], allowed, schema) + ordinary = {"type": "array", "items": {"type": "integer"}} + assert _bfcl_argument_matches([1, 2], [[1, 2]], ordinary) + assert not _bfcl_argument_matches([1, 2], [1, 2], ordinary) diff --git a/tests/kernels/attention/test_qsa_benchmark_indices.py b/tests/kernels/attention/test_qsa_benchmark_indices.py new file mode 100644 index 0000000000..a690af5a55 --- /dev/null +++ b/tests/kernels/attention/test_qsa_benchmark_indices.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU metadata checks; the benchmark import requires a Flash-V100 installation.""" + +import pytest +import torch + +pytest.importorskip("flash_attn_v100") + +from benchmarks.kernels.benchmark_qwen38_qsa_mtp5_flash_v100 import _logical_indices + + +@pytest.mark.parametrize("rows", (4, 5, 8, 16)) +@pytest.mark.parametrize("independent", (False, True)) +@pytest.mark.parametrize("overlap", (0.0, 0.82, 1.0)) +def test_qsa_micro_uses_canonical_causal_selections(rows, independent, overlap): + indices = _logical_indices( + rows=rows, + seq_len=8192, + overlap=overlap, + seed=7, + independent=independent, + ) + assert indices.shape == (rows, 2051) + for row in range(rows): + visible = 8192 if independent else 8192 - rows + row + 1 + complete = indices[row, :2048].reshape(512, 4) + assert torch.all(complete[:, 0] % 4 == 0) + assert torch.all(complete[:, 0] // 4 < visible // 4) + torch.testing.assert_close( + complete, + complete[:, :1] + torch.arange(4, dtype=indices.dtype), + ) + tail = indices[row, 2048:] + count = visible % 4 + torch.testing.assert_close( + tail[:count], + torch.arange(visible - count, visible, dtype=indices.dtype), + ) + assert torch.all(tail[count:] == -1) + valid = indices[row][indices[row] >= 0] + assert valid.unique().numel() == 2048 + count + assert torch.all(valid < visible) diff --git a/tests/kernels/attention/test_sm70_flashinfer_mqa.py b/tests/kernels/attention/test_sm70_flashinfer_mqa.py new file mode 100644 index 0000000000..03606cdd06 --- /dev/null +++ b/tests/kernels/attention/test_sm70_flashinfer_mqa.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native scheduler/scorer tests. These are not model-quality admission.""" + +import pytest +import torch + +from benchmarks.kernels.benchmark_sm70_flashinfer_mqa import ( + check_schedule, + make_inputs, +) +from benchmarks.kernels.flashinfer_sm70_mqa import FlashInferMQA, build + +pytestmark = [ + pytest.mark.skip_global_cleanup, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +@pytest.fixture(scope="module", autouse=True) +def native(): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("Native Volta test") + build() + + +def oracle(inputs, op): + q, k, table, requests, positions, lengths = inputs + expected_visible = [] + for row in range(q.shape[0]): + req, pos = int(requests[row]), int(positions[row]) + n = 0 + if 0 <= req < len(lengths) and pos >= 0: + n = max( + 0, + min( + (pos + 1) // 4, + int(lengths[req]) // 4, + op.logits.shape[1], + table.shape[1] * k.shape[1], + ), + ) + expected_visible.append(n) + if n: + columns = torch.arange(n, device=q.device) + pages = table[req, columns // k.shape[1]].long() + live = (pages >= 0) & (pages < k.shape[0]) + keys = k[pages.clamp(0, k.shape[0] - 1), columns % k.shape[1], 0] + scores = (keys.double() @ q[row].double().T).clamp_min(0).sum(-1) + scores /= q.shape[2] ** 0.5 + scores[~live] = -torch.inf + torch.testing.assert_close( + op.logits[row, :n].double(), scores, rtol=2e-5, atol=2e-5 + ) + assert torch.isnan(op.logits[row, n:]).all(), "Unowned tail was overwritten" + assert op.visible.cpu().tolist() == expected_visible + check_schedule(op) + + +@pytest.mark.parametrize("rows", [1, 2, 4, 8, 16, 32, 33, 64]) +@pytest.mark.parametrize("dim,heads", [(64, 1), (128, 4), (256, 16)]) +def test_graph_dynamic_lengths_empty_rows_and_invalid_pages(rows, dim, heads): + inputs = make_inputs(rows, 1024, dim=dim, heads=heads, table_width=3) + q, k, table, requests, positions, lengths = inputs + op = FlashInferMQA(q, 581, 80) + op(*inputs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*inputs) + for iteration in range(4): + lengths.copy_( + torch.tensor( + [0 if i % 3 == iteration else 1024 + i * 7 for i in range(rows)], + dtype=torch.int32, + device=q.device, + ) + ) + positions.copy_(lengths - 1) + if iteration == 1: + requests[0] = -1 + elif iteration == 2: + requests[0] = rows + positions[-1] = -19 + else: + requests.copy_(torch.arange(rows, device=q.device)) + table[:, 0] = -1 if iteration == 1 else k.shape[0] if iteration == 2 else 0 + op.logits.fill_(torch.nan) + op.visible.fill_(-993) + op.schedule.fill_(-991) + q.mul_(-0.875) + graph.replay() + oracle(inputs, op) + + +def test_distinct_graph_instances_and_strided_inputs(): + inputs = make_inputs(8, 2048, table_width=3) + q, k, table, requests, positions, lengths = inputs + # Valid noncompact leading strides, including separate key/value planes. + q_pad = torch.zeros(8, 4, 256, dtype=q.dtype, device=q.device) + q_pad[..., :128].copy_(q) + k_pad = torch.zeros(k.shape[0], 2, *k.shape[1:], device=k.device, dtype=k.dtype) + k_pad[:, 0].copy_(k) + inputs = q_pad[..., :128], k_pad[:, 0], table, requests, positions, lengths + ops = [FlashInferMQA(inputs[0], 588, 80) for _ in range(2)] + streams = [torch.cuda.Stream() for _ in ops] + graphs = [torch.cuda.CUDAGraph() for _ in ops] + torch.accelerator.synchronize() + for stream, graph, op in zip(streams, graphs, ops): + with torch.cuda.stream(stream): + op(*inputs) + stream.synchronize() + with torch.cuda.graph(graph, stream=stream): + op(*inputs) + for op in ops: + op.logits.fill_(torch.nan) + torch.accelerator.synchronize() + for _ in range(7): + for stream, graph in zip(streams, graphs): + with torch.cuda.stream(stream): + graph.replay() + torch.accelerator.synchronize() + for op in ops: + oracle(inputs, op) + torch.testing.assert_close( + ops[0].logits, ops[1].logits, equal_nan=True, atol=0, rtol=0 + ) + + +def test_metadata_guard_rejects_wrong_dtype_before_launch(): + inputs = make_inputs(4, 1024, table_width=3) + op = FlashInferMQA(inputs[0], 588, 80) + args = list(inputs) + args[3] = args[3].long() + with pytest.raises(RuntimeError, match="int32"): + op(*args) + + +def test_int32_and_int64_positions_agree_and_large_position_does_not_wrap(): + inputs = list(make_inputs(4, 1024, table_width=3)) + op = FlashInferMQA(inputs[0], 588, 80) + op(*inputs) + expected = op.logits.clone() + inputs[4] = inputs[4].int() + op(*inputs) + torch.testing.assert_close(op.logits[:, :256], expected[:, :256], atol=0, rtol=0) + inputs[4] = torch.full_like( + inputs[4], torch.iinfo(torch.int64).max, dtype=torch.int64 + ) + op.logits.fill_(torch.nan) + op(*inputs) + oracle(inputs, op) + + +@pytest.mark.parametrize("rows,workers", [(1, 80), (4, 80), (16, 320), (64, 160)]) +def test_every_live_tile_is_written_by_exactly_one_cta(rows, workers): + inputs = make_inputs(rows, 2048, table_width=3) + q, k, table, requests, positions, lengths = inputs + op = FlashInferMQA(q, 588, workers) + visits = torch.empty((rows, 10), device=q.device, dtype=torch.int32) + for repeat in range(3): + lengths.copy_( + torch.tensor( + [0 if row % 3 == repeat else 255 + 197 * row for row in range(rows)], + device=q.device, + dtype=torch.int32, + ) + ) + positions.copy_(lengths - 1) + visits.zero_() + torch.ops._C_flashinfer_mqa_sm70.run( + *inputs, op.logits, op.visible, op.schedule, 4, 128**0.5, workers, visits + ) + expected = ( + torch.arange(10, device=q.device)[None] < ((op.visible + 63) // 64)[:, None] + ) + assert torch.equal(visits, expected.int()) diff --git a/tests/kernels/test_hc_benchmark_weight_views.py b/tests/kernels/test_hc_benchmark_weight_views.py new file mode 100644 index 0000000000..60a9c0e1ac --- /dev/null +++ b/tests/kernels/test_hc_benchmark_weight_views.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU layout proof only; CUDA performance/rounding require the TP4 screen.""" + +import pytest +import torch + +from benchmarks.kernels.benchmark_sm70_hc_batch_tp4 import shard_weights, up_projection + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.mark.parametrize("rank", range(4)) +@pytest.mark.parametrize("rows", (1, 4, 8, 16)) +def test_hc_weight_views_alias_original_without_output_overlap(rank, rows): + torch.manual_seed(7) + # Integral low-magnitude inputs make the layout oracle exact on CPU; + # this deliberately does not purport to test GPU reduction association. + down = torch.randint(-1, 2, (336, 10240)).half() + up = torch.randint(-1, 2, (10240, 320)).half() + lora = torch.randint(-1, 2, (rows, 320)).half() + wd, wu = shard_weights(down, up, rank, "views") + pd, pu = shard_weights(down, up, rank, "packed") + assert wd.untyped_storage().data_ptr() == down.untyped_storage().data_ptr() + assert wu.untyped_storage().data_ptr() == up.untyped_storage().data_ptr() + assert wd.storage_offset() == rank * 80 * 10240 + torch.testing.assert_close(wd[:80], pd[:80], rtol=0, atol=0) + if rank == 3: + torch.testing.assert_close(wd[80:84], pd[80:84], rtol=0, atol=0) + output = torch.full((rows, 2560), float("nan"), dtype=torch.float16) + result = up_projection(lora, wu, output) + expected = torch.nn.functional.linear(lora, pu) + assert result.data_ptr() == output.data_ptr() + torch.testing.assert_close(result, expected, rtol=0, atol=0) + offsets = torch.arange(rows * 2560).view(rows, 4, 640).transpose(0, 1) + assert offsets.unique().numel() == rows * 2560 + # Refreshing weights must be visible without stale packed copies. + up.zero_() + up_projection(lora, wu, output) + assert torch.count_nonzero(output).item() == 0 diff --git a/tests/kernels/test_sm70_flashinfer_gdn_native.py b/tests/kernels/test_sm70_flashinfer_gdn_native.py new file mode 100644 index 0000000000..2ff7cf8774 --- /dev/null +++ b/tests/kernels/test_sm70_flashinfer_gdn_native.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Wheel-native GDN geometry isolation and replay/state ownership checks. + +TP1/2/4 head partitions of the same layer must produce the same result. This +is an operator test, not a model-score or distributed-communication gate. +""" + +import importlib.util +import os + +import pytest +import torch + + +@pytest.fixture(scope="module") +def native(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + path = os.environ.get("SM70_FLASHINFER_GDN_TEST_LIBRARY") + if not path: + spec = importlib.util.find_spec("vllm._sm70_flashinfer_gdn_C") + path = spec.origin if spec else None + if not path: + pytest.skip("Install the native GDN wheel fragment or set its test path") + torch.ops.load_library(path) + + +@pytest.mark.parametrize("rows", [1, 2, 4, 8, 16, 32, 64]) +def test_head_partition_graph_state_equivalence(native, rows): + torch.manual_seed(20260906) + device = "cuda" + pool, hq, hv, dim, hidden = rows + 3, 16, 48, 128, 2560 + width = (2 * hq + hv) * dim + x = torch.randn(rows, hidden, device=device, dtype=torch.float16) * 0.1 + qkv = torch.randn(rows, width, device=device, dtype=torch.float16) * 0.1 + weights = torch.randn(hidden, 2 * hv, device=device, dtype=torch.float16) * 0.01 + cw = torch.randn(width, 4, device=device, dtype=torch.float16) * 0.1 + bias = torch.randn(width, device=device, dtype=torch.float16) * 0.01 + a = torch.randn(hv, device=device, dtype=torch.float32) * 0.1 + dt = torch.randn(hv, device=device, dtype=torch.float16) * 0.1 + conv = torch.randn(pool, width, 3, device=device, dtype=torch.float16) * 0.1 + state = torch.randn(pool, hv, dim, dim, device=device) * 0.01 + indices = torch.arange(rows, device=device, dtype=torch.int32) + groups = [] + for tp in (1, 2, 4): + parts = [] + q, v = hq // tp, hv // tp + for rank in range(tp): + qs = torch.arange(rank * q * dim, (rank + 1) * q * dim, device=device) + vs = torch.arange(rank * v * dim, (rank + 1) * v * dim, device=device) + channels = torch.cat((qs, qs + hq * dim, vs + 2 * hq * dim)) + heads = torch.arange(rank * v, (rank + 1) * v, device=device) + columns = torch.cat((heads, heads + hv)) + # Use production's SD storage represented as a [pool, C, 3] view. + c = conv[:, channels].transpose(1, 2).contiguous().transpose(1, 2) + # Preserve a padded pool stride rather than requiring dense states. + s = torch.empty(pool * 2, v, dim, dim, device=device)[::2] + s.copy_(state[:, heads]) + raw = qkv[:, channels].contiguous() + out = torch.empty(rows, v, dim, device=device, dtype=torch.float16) + conv_out = torch.empty_like(raw) + partial = torch.empty(rows * 2 * v * 160, device=device) + run = getattr(torch.ops, f"_C_flashinfer_gdn_sm70_h2560_q{q}_v{v}").run + args = ( + x, + weights[:, columns].contiguous(), + raw, + cw[channels], + bias[channels], + c, + a[heads], + dt[heads], + s, + indices, + out, + conv_out, + partial, + ) + parts.append( + dict( + run=run, + args=args, + raw=raw, + channels=channels, + heads=heads, + state=s, + conv=c, + out=out, + ) + ) + groups.append(parts) + + def launch(): + for parts in groups: + for p in parts: + p["run"](*p["args"]) + + for _ in range(3): + launch() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + for parts in groups: + for p in parts: + p["conv"].copy_(conv[:, p["channels"]]) + p["state"].copy_(state[:, p["heads"]]) + for step in range(8): + x.normal_().mul_(0.1) + qkv.normal_().mul_(0.1) + indices.copy_(torch.randperm(pool, device=device)[:rows]) + if step % 3 == 1: + indices[-1] = -1 + for parts in groups: + for p in parts: + p["raw"].copy_(qkv[:, p["channels"]]) + p["out"].fill_(torch.nan) + graph.replay() + full = groups[0][0] + for parts in groups[1:]: + for p in parts: + torch.testing.assert_close( + p["conv"], full["conv"][:, p["channels"]], atol=0, rtol=0 + ) + torch.testing.assert_close( + p["state"], full["state"][:, p["heads"]], atol=0, rtol=0 + ) + torch.testing.assert_close( + p["out"], full["out"][:, p["heads"]], atol=0, rtol=0 + ) + assert torch.isfinite(p["out"]).all() diff --git a/tests/kernels/test_sm70_qsa_page4_plan.py b/tests/kernels/test_sm70_qsa_page4_plan.py new file mode 100644 index 0000000000..453ed159c1 --- /dev/null +++ b/tests/kernels/test_sm70_qsa_page4_plan.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Allocation-invariance contract for the SM70 grouped QSA planner. + +Keep logical selections, query grouping, and physical-page aliasing fixed. +Relocating KV pages must preserve masks, category padding, and the logical +reduction order. This is not a cross-batch-shape invariance contract. +""" + +import pytest +import torch + +WIDTH = 2051 +OUTPUT_WIDTH = 4160 + + +@pytest.fixture +def extension(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("requires a V100 / SM70 GPU") + return pytest.importorskip("flash_attn_v100_cuda") + + +def make_case(kind="mixed", page_size=16, permute_selection=False): + generator = torch.Generator().manual_seed(387466) + if kind == "wide": + requests = torch.arange(8, dtype=torch.int32) + lengths = torch.full((8,), 8191, dtype=torch.int32) + positions = lengths.to(torch.int64) - 1 + else: + requests = torch.tensor([0, 0, 1, 1, 2, 3, 3, 3], dtype=torch.int32) + lengths = torch.tensor([61, 57, 59, 63], dtype=torch.int32) + positions = torch.tensor([16, 38, 19, 40, 36, 25, 32, 62]) + pages = (int(lengths.max()) + page_size - 1) // page_size + table = torch.arange(len(lengths) * pages, dtype=torch.int32).view(-1, pages) + if kind == "shared": + table[:, 0] = table[0, 0] + indices = torch.full((8, WIDTH), -1, dtype=torch.int32) + for row, request in enumerate(requests.tolist()): + visible = min(int(positions[row]) + 1, int(lengths[request])) + blocks = torch.randperm(visible // 4, generator=generator)[:512].sort().values + if permute_selection: + blocks = blocks.flip(0) + full = (blocks[:, None] * 4 + torch.arange(4)).flatten() + indices[row, : len(full)] = full.to(torch.int32) + tail = torch.arange(visible // 4 * 4, visible, dtype=torch.int32) + indices[row, len(full) : len(full) + len(tail)] = tail + if kind == "invalid": + requests[1], requests[4] = -1, len(lengths) + indices[6] = -1 + if kind == "empty": + indices.fill_(-1) + return indices, table, requests, positions, lengths + + +def relocate(case, layout): + indices, table, requests, positions, lengths = case + count = int(table.max()) + 1 + if layout == "collision": + # PAGE16, interleaved K/V: stride=8, so 1024 cache pages collide + # in the 8192-entry hash. No K/V allocation is needed for plan tests. + mapping = torch.arange(count, dtype=torch.int32) * 2048 + 7 + elif layout == "shuffled": + mapping = torch.randperm(count, generator=torch.Generator().manual_seed(466)) + mapping = mapping.to(torch.int32) * 3 + 11 + else: + mapping = torch.arange(count, dtype=torch.int32) + return (indices, mapping[table.long()], requests, positions, lengths), mapping + + +def category(mask): + result = 0 + for query in range(8): + if mask & (15 << (query * 4)): + result |= 1 << (query * 6 // 16) + result |= 1 << ((query * 6 + 5) // 16) + return result + + +def reference(case, page_size, physical_stride): + indices, table, requests, positions, lengths = case + entries: dict[int, tuple[int, int]] = {} + for query, request in enumerate(requests.tolist()): + if not 0 <= request < len(lengths): + continue + visible = min(max(int(positions[query]) + 1, 0), int(lengths[request])) + count = min(visible // 4, 512) * 4 + visible % 4 + for token in indices[query, :count].tolist(): + if not 0 <= token < visible: + continue + physical = int(table[request, token // page_size]) + physical = physical * physical_stride + token % page_size // 4 + owner = (query << 29) | (token // 4) + old_mask, old_owner = entries.get(physical, (0, (1 << 32) - 1)) + entries[physical] = ( + old_mask | (1 << (query * 4 + token % 4)), + min(old_owner, owner), + ) + pages: list[int] = [] + masks: list[int] = [] + for group_category in range(1, 8): + bucket = sorted( + (owner, physical, mask) + for physical, (mask, owner) in entries.items() + if category(mask) == group_category + ) + pages.extend(physical for _, physical, _ in bucket) + masks.extend(mask for _, _, mask in bucket) + padding = -len(bucket) % 8 + pages.extend([0] * padding) + masks.extend([0] * padding) + return pages, masks + + +def run_plan(extension, case, page_size, physical_stride, num_cache_blocks=None): + device_case = tuple(t.cuda() for t in case) + if num_cache_blocks is None: + num_cache_blocks = int(case[1].max()) + 1 + pages = torch.full((1, OUTPUT_WIDTH), -17, dtype=torch.int32, device="cuda") + masks = torch.zeros_like(pages, dtype=torch.uint32) + lengths = torch.empty(1, dtype=torch.int32, device="cuda") + + def launch(): + extension.grouped_sparse_page4_plan_fwd( + *device_case, + pages, + masks, + lengths, + page_size, + physical_stride, + num_cache_blocks, + ) + + launch() + return pages, masks, lengths, launch, device_case + + +@pytest.mark.parametrize("kind", ["mixed", "shared", "invalid", "empty", "wide"]) +@pytest.mark.parametrize("layout", ["compact", "shuffled", "collision"]) +def test_plan_matches_logical_reference(extension, kind, layout): + case, _ = relocate(make_case(kind), layout) + expected_pages, expected_masks = reference(case, 16, 8) + pages, masks, lengths, launch, _ = run_plan(extension, case, 16, 8) + assert int(lengths[0]) == len(expected_pages) * 4 + count = len(expected_pages) + assert pages[0, :count].cpu().tolist() == expected_pages + assert masks[0, :count].cpu().tolist() == expected_masks + first = (pages.clone(), masks.clone(), lengths.clone()) + for _ in range(3): + launch() + for before, after in zip(first, (pages, masks, lengths)): + assert torch.equal(before, after) + + +@pytest.mark.parametrize("page_size", [4, 16, 32]) +def test_plan_selection_order_and_graph_relocation(extension, page_size): + case = make_case("shared", page_size) + changed_case = make_case("shared", page_size, permute_selection=True) + changed_case, _ = relocate(changed_case, "shuffled") + stride = page_size // 4 + pages, masks, lengths, launch, device_case = run_plan( + extension, case, page_size, stride, int(changed_case[1].max()) + 1 + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + for source, target in zip(changed_case, device_case): + target.copy_(source) + expected_pages, expected_masks = reference(changed_case, page_size, stride) + for _ in range(3): + graph.replay() + count = len(expected_pages) + assert int(lengths[0]) == count * 4 + assert pages[0, :count].cpu().tolist() == expected_pages + assert masks[0, :count].cpu().tolist() == expected_masks + + +@pytest.mark.parametrize("reverse", [False, True]) +def test_plan_live_union_crosses_sort_tiles_during_graph_replay(extension, reverse): + # Eight independent requests can contribute 512 complete microblocks and + # one partial tail each. Exercise every compact-sort transition, including + # empty -> maximum -> empty, without recapturing or changing addresses. + sizes = [ + 0, + 1, + 511, + 512, + 513, + 1023, + 1024, + 1025, + 2047, + 2048, + 2049, + 4095, + 4096, + 4097, + 4104, + 0, + ] + if reverse: + sizes.reverse() + table = torch.arange(8 * 513, dtype=torch.int32).reshape(8, 513) + requests = torch.arange(8, dtype=torch.int32) + + def make_union(size): + counts = torch.tensor([size // 8 + (row < size % 8) for row in range(8)]) + visible = torch.where(counts == 513, 2049, counts * 4).to(torch.int32) + indices = torch.arange(WIDTH, dtype=torch.int32).expand(8, -1).clone() + indices[indices >= visible[:, None]] = -1 + return indices, table, requests, visible.to(torch.int64) - 1, visible + + pages, masks, lengths, launch, device_case = run_plan( + extension, make_union(0), 4, 1 + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + for size in sizes: + case = make_union(size) + for source, target in zip(case, device_case): + target.copy_(source) + expected_pages, expected_masks = reference(case, 4, 1) + for _ in range(2): + graph.replay() + count = len(expected_pages) + assert int(lengths[0]) == count * 4 + assert pages[0, :count].cpu().tolist() == expected_pages + assert masks[0, :count].cpu().tolist() == expected_masks + + +@pytest.mark.parametrize("interleaved", [False, True]) +@pytest.mark.parametrize("kv_dtype", ["auto", "fp8_e4m3"]) +def test_attention_is_bitwise_invariant_to_physical_relocation( + extension, interleaved, kv_dtype +): + from vllm.models.qwen4_exp.nvidia.ops import qsa + + case = make_case("shared") + generator = torch.Generator(device="cuda").manual_seed(387) + count = int(case[1].max()) + 1 + kv = torch.randn( + count, 2, 16, 1, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + if kv_dtype == "fp8_e4m3": + kv = kv.to(torch.float8_e4m3fn).view(torch.uint8) + query = torch.randn( + 8, 6, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + outputs = [] + for layout in ("compact", "shuffled"): + remapped, mapping = relocate(case, layout) + cache = torch.zeros( + int(mapping.max()) + 1, 2, 16, 1, 256, dtype=kv.dtype, device="cuda" + ) + cache[mapping.long().cuda()] = kv + key, value = cache[:, 0], cache[:, 1] + if not interleaved: + key, value = key.contiguous(), value.contiguous() + stride = key.stride(0) // (4 * 256) + pages, masks, lengths, _, _ = run_plan(extension, remapped, 16, stride) + physical_k, physical_v = qsa._qsa_xqa_page4_physical_kv(query, key, value) + out = torch.empty_like(query) + lse = torch.empty((8, 6), dtype=torch.float32, device="cuda") + extension.grouped_sparse_page4_fwd( + query, + physical_k, + physical_v, + out, + pages, + masks, + lengths, + lse, + 256**-0.5, + kv_dtype, + 0.125, + 0.25, + ) + outputs.append(out.clone()) + assert torch.equal(*outputs) + + +@pytest.mark.parametrize("permute_selection", [False, True]) +def test_single_row_page4_table_uses_logical_order(extension, permute_selection): + from vllm.models.qwen4_exp.nvidia.ops import qsa + + case, _ = relocate(make_case(permute_selection=permute_selection), "shuffled") + indices, table, requests, positions, lengths = case + pages, seq_lens = qsa._qsa_xqa_page4_block_table( + *(tensor.cuda() for tensor in case), int(table.max()) + 1, 16, 8 + ) + for row, request in enumerate(requests.tolist()): + visible = min(int(positions[row]) + 1, int(lengths[request])) + tokens = sorted(indices[row, :visible:4].tolist()) + expected = [int(table[request, t // 16]) * 8 + t % 16 // 4 for t in tokens] + assert pages[row, : len(expected)].cpu().tolist() == expected + assert int(seq_lens[row]) == visible + + +@pytest.mark.parametrize("kv_dtype", ["auto", "fp8_e4m3"]) +def test_mixed_grouped_and_xqa_tail_is_allocation_invariant( + extension, monkeypatch, kv_dtype +): + from vllm.models.qwen4_exp.nvidia.ops import qsa + + monkeypatch.setattr(qsa, "_SM70_QSA_XQA_PAGE4", True) + monkeypatch.setattr(qsa, "_SM70_QSA_XQA_PAGE4_MIN_ROWS", 64) + monkeypatch.setattr(qsa, "_SM70_QSA_GROUPED_PAGE4", True) + calls = [] + grouped = qsa._qsa_sparse_paged_attention_sm70_grouped_page4 + tail = qsa._qsa_sparse_paged_attention_sm70_xqa_page4_batch + + def record_grouped(query, *args): + calls.append(("grouped", query.shape[0])) + return grouped(query, *args) + + def record_tail(query, *args): + calls.append(("tail", query.shape[0])) + return tail(query, *args) + + monkeypatch.setattr( + qsa, "_qsa_sparse_paged_attention_sm70_grouped_page4", record_grouped + ) + monkeypatch.setattr( + qsa, "_qsa_sparse_paged_attention_sm70_xqa_page4_batch", record_tail + ) + case = make_case("shared") + generator = torch.Generator(device="cuda").manual_seed(185) + count = int(case[1].max()) + 1 + kv = torch.randn( + count, 2, 16, 1, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + if kv_dtype == "fp8_e4m3": + kv = kv.to(torch.float8_e4m3fn).view(torch.uint8) + query = torch.randn( + 185, 6, 256, generator=generator, dtype=torch.float16, device="cuda" + ) + outputs = [] + for layout in ("compact", "shuffled"): + remapped, mapping = relocate(case, layout) + indices, table, requests, positions, lengths = [t.cuda() for t in remapped] + indices = indices.repeat(24, 1)[:185].contiguous() + requests = requests.repeat(24)[:185].contiguous() + positions = positions.repeat(24)[:185].contiguous() + cache = torch.zeros( + int(mapping.max()) + 1, 2, 16, 1, 256, dtype=kv.dtype, device="cuda" + ) + cache[mapping.long().cuda()] = kv + outputs.append( + qsa.qsa_sparse_paged_attention( + query, + cache[:, 0], + cache[:, 1], + indices, + table, + requests, + query_positions=positions, + sequence_lengths=lengths, + kv_cache_dtype=kv_dtype, + k_scale=0.125, + v_scale=0.25, + ).clone() + ) + assert torch.equal(*outputs) + assert calls == [("grouped", 184), ("tail", 1)] * 2 diff --git a/tests/kernels/test_sm70_qsa_rounding_isolation.py b/tests/kernels/test_sm70_qsa_rounding_isolation.py new file mode 100644 index 0000000000..e26ae18c52 --- /dev/null +++ b/tests/kernels/test_sm70_qsa_rounding_isolation.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU sanity gates for the numerical counterfactual, not native CUDA gates.""" + +import pytest +import torch + +from benchmarks.kernels.benchmark_sm70_qsa_rounding_isolation import isolate + + +@pytest.mark.parametrize("rounded", [False, True]) +def test_order_duplicates_padding_and_masked_nan(rounded): + # Zero logits make every valid probability exactly one before normalization. + k = torch.zeros(2, 16, 1, 16, dtype=torch.float16) + v = torch.arange(32).half().reshape(2, 16, 1, 1).expand_as(k).clone() + k[0, 0] = torch.nan + v[0, 0] = torch.nan + capture = dict( + q=torch.zeros(2, 2, 16, dtype=torch.float16), + k=k, + v=v, + indices=torch.tensor([[18, 2, 18, -1] + [-1] * 13] * 2), + table=torch.tensor([[0, 1], [0, 1]]), + requests=torch.tensor([0, -1]), + gate=None, + ) + result = isolate(capture, rounded) + expected = torch.zeros_like(result) + expected[0].fill_(38 / 3) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + assert torch.isfinite(result).all() + + +def test_probability_cast_does_not_round_the_denominator(): + q = torch.ones(1, 1, 16, dtype=torch.float16) + k = torch.zeros(1, 16, 1, 16, dtype=torch.float16) + k[0, 1] = 0.31 + v = torch.zeros_like(k) + v[0, 0] = 1 + v[0, 1] = -0.8 + capture = dict( + q=q, + k=k, + v=v, + indices=torch.tensor([[0, 1]]), + table=torch.tensor([[0]]), + requests=torch.tensor([0]), + gate=None, + ) + scores = (q.double()[:, 0] @ k[0, :2, 0].double().t())[0] / 4 + p = (scores - scores.max()).exp() + expected = (p.half().double() @ v[0, :2, 0].double()) / p.sum() + torch.testing.assert_close( + isolate(capture, True)[0, 0], expected.half(), atol=0, rtol=0 + ) diff --git a/tests/models/qwen4_exp/test_sm70_batch_hc.py b/tests/models/qwen4_exp/test_sm70_batch_hc.py new file mode 100644 index 0000000000..d4d66e7c8d --- /dev/null +++ b/tests/models/qwen4_exp/test_sm70_batch_hc.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace as NS + +import pytest +import torch +from torch import nn + +from vllm import _custom_ops as ops +from vllm import envs +from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator +from vllm.forward_context import is_uniform_decode_metadata +from vllm.models.qwen4_exp.nvidia import sm70_batch_hc as hc + +pytestmark = pytest.mark.skip_global_cleanup + + +def test_batch_hc_is_opt_in(monkeypatch): + envs.disable_envs_cache() + monkeypatch.delenv("VLLM_SM70_QWEN38_BATCH_HC_FP16", raising=False) + assert not envs.VLLM_SM70_QWEN38_BATCH_HC_FP16 + + +@pytest.mark.parametrize("rows", (1, 2, 4, 8, 16, 32)) +def test_cpu_metadata_is_independent_of_scheduler_and_kv(rows): + assert is_uniform_decode_metadata( + { + "attn": NS(max_query_len=1), + "gdn": NS( + num_prefills=0, + num_decodes=rows, + num_decode_tokens=rows, + ), + } + ) + + +@pytest.mark.parametrize( + "metadata", + ( + None, + {}, + [], + {"unknown": NS()}, + {"attn": NS(max_query_len=5)}, + {"attn": NS(max_query_len=1), "gdn": NS(num_prefills=1)}, + {"gdn": NS(num_decodes=4, num_decode_tokens=8)}, + {"attn": NS(max_query_len=torch.tensor(1))}, + ), +) +def test_mixed_prefill_verify_unknown_context_falls_back(metadata): + assert not is_uniform_decode_metadata(metadata) + + +@pytest.mark.parametrize("rows", (1, 17, 32)) +def test_unsupported_width_delegates_to_original_fused_hc(monkeypatch, rows): + x = torch.empty(rows, 10240, dtype=torch.float16) + expected = (torch.empty(rows, 2560), torch.empty(rows, 4)) + monkeypatch.setattr(hc, "_channel", lambda: pytest.fail("unexpected batch channel")) + monkeypatch.setattr(hc, "_qwen38_sm70_fp16_fused_hc", lambda *a: expected) + dummy = torch.empty(0) + assert hc._batch_hc(x, dummy, dummy, dummy, True, False) is expected + + +def test_original_gemv_fallback_is_not_replaced_by_linear(monkeypatch): + x = torch.empty(1, 10240, dtype=torch.float16) + seen = [] + + def gemv(x, weight): + seen.append("gemv") + return torch.zeros(1, 336, dtype=x.dtype) + + monkeypatch.setattr(hc, "_qwen38_sm70_fp16_gemv", gemv) + monkeypatch.setattr(hc, "hc_silu", lambda a, n: a) + monkeypatch.setattr(hc, "hc_gate_mix", lambda a, b, n: b[:, :2560]) + hc._batch_hc( + x, + torch.empty(0), + torch.zeros(10240, 320, dtype=x.dtype), + torch.empty(0), + False, + True, + ) + assert seen == ["gemv"] + + +@pytest.mark.parametrize("rows", (2, 4, 8, 16)) +@pytest.mark.parametrize("rank", range(4)) +def test_batch_hc_preserves_full_down_and_only_shards_up(monkeypatch, rows, rank): + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0") + envs.disable_envs_cache() + x = torch.empty(rows, 10240, dtype=torch.float16) + down = torch.empty(336, 10240, dtype=x.dtype) + up = torch.empty(10240, 320, dtype=x.dtype) + packed = torch.empty(2560, 320, dtype=x.dtype) + projected = torch.arange(rows * 336, dtype=torch.float32).reshape(rows, 336) + projected = projected.to(x.dtype) + lora = torch.empty(rows, 320, dtype=x.dtype) + gate = torch.empty(rows, 2560, dtype=x.dtype) + calls = [] + + def linear(value, weight): + if weight is down: + assert value is x + calls.append("full_down") + return projected + assert weight is packed and value is lora + calls.append("sharded_up") + return gate + + def silu(value, groups): + assert groups == 4 + torch.testing.assert_close(value, projected[:, :320]) + calls.append("silu") + return lora + + def mix(ptr, actual_gate, actual_x, block): + assert ptr == 42 and actual_gate is gate and actual_x is x + calls.append("mix") + block.zero_() + + monkeypatch.setattr(hc, "_decode_context_ok", lambda: True) + monkeypatch.setattr( + hc, + "_channel", + lambda: NS(rank=rank, _ptr=42, can_sm70_qwen38_hc_batch=lambda x: True), + ) + monkeypatch.setattr(torch.nn.functional, "linear", linear) + monkeypatch.setattr(hc, "hc_silu", silu) + monkeypatch.setattr( + ops, + "sm70_qwen38_hc_batch_down", + lambda *a: pytest.fail("down must not be sharded or gathered"), + ) + monkeypatch.setattr(ops, "sm70_qwen38_hc_batch_mix", mix) + block, injection = hc._batch_hc(x, down, up, packed, True, False) + assert calls == ["full_down", "silu", "sharded_up", "mix"] + assert block.shape == (rows, 2560) + assert injection.is_contiguous() + torch.testing.assert_close(injection, projected[:, 320:324], rtol=0, atol=0) + + +@pytest.mark.parametrize("rank", range(4)) +def test_up_shard_reload_keeps_captured_pointer_and_is_nonpersistent(rank): + child = nn.Module() + weight = torch.arange(10240, dtype=torch.float16)[:, None].expand(-1, 320).clone() + hc._copy_up_shard(child, weight, rank) + pointer = child._sm70_batch_hc_up.data_ptr() + assert not child.state_dict() + expected = weight.view(4, 2560, 320)[:, rank * 640 : (rank + 1) * 640] + torch.testing.assert_close(child._sm70_batch_hc_up, expected.reshape(2560, 320)) + weight.zero_() + hc._copy_up_shard(child, weight, rank) + assert child._sm70_batch_hc_up.data_ptr() == pointer + assert torch.count_nonzero(child._sm70_batch_hc_up).item() == 0 + + +def test_channel_destroy_is_idempotent_and_does_not_close_ordinary_channel(): + closed = [] + communicator = object.__new__(CudaCommunicator) + communicator.sm70_hc_batch_comm = NS(close=lambda: closed.append("hc")) + communicator.pynccl_comm = communicator.ca_comm = None + communicator.fi_ar_comm = communicator.all2all_manager = None + communicator.destroy() + communicator.destroy() + assert closed == ["hc"] + + +def test_native_capabilities_resolve_only_from_the_pointer_owner(monkeypatch): + monkeypatch.setattr(ops, "_custom_ar_owner_namespace", lambda: NS()) + assert not ops.supports_sm70_qwen38_hc_batch() + called = [] + owner = NS( + sm70_qwen38_hc_batch_down=lambda *a: called.append("down"), + sm70_qwen38_hc_batch_mix=lambda *a: called.append("mix"), + ) + monkeypatch.setattr(ops, "_custom_ar_owner_namespace", lambda: owner) + assert ops.supports_sm70_qwen38_hc_batch() + dummy = torch.empty(0) + ops.sm70_qwen38_hc_batch_down(1, dummy, dummy, dummy) + ops.sm70_qwen38_hc_batch_mix(1, dummy, dummy, dummy) + assert called == ["down", "mix"] + + +def test_fake_hc_never_touches_communicator(monkeypatch): + from torch._subclasses.fake_tensor import FakeTensorMode + + monkeypatch.setattr(hc, "_channel", lambda: pytest.fail("fake cannot use channel")) + with FakeTensorMode(): + x = torch.empty(8, 10240, dtype=torch.float16) + block, injection = torch.ops.vllm.qwen38_sm70_batch_hc( + x, + torch.empty(336, 10240), + torch.empty(10240, 320), + torch.empty(2560, 320), + True, + False, + ) + assert block.shape == (8, 2560) and injection.shape == (8, 4) + + +def test_derived_weight_hook_runs_after_quant_postprocessing(monkeypatch): + from vllm.model_executor.model_loader import utils + + events = [] + + class Quant: + def process_weights_after_loading(self, layer): + events.append("quant") + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.quant_method = Quant() + + def prepare_sm70_batch_hc(self): + assert events == ["quant"] + events.append("hc") + + monkeypatch.setattr(utils, "QuantizeMethodBase", Quant) + utils.process_weights_after_loading( + Model(), + NS(dtype=torch.float16, quantization=None), + torch.device("cpu"), + ) + assert events == ["quant", "hc"] diff --git a/tests/models/qwen4_exp/test_sm70_gdn_projection_split.py b/tests/models/qwen4_exp/test_sm70_gdn_projection_split.py new file mode 100644 index 0000000000..140931533b --- /dev/null +++ b/tests/models/qwen4_exp/test_sm70_gdn_projection_split.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.envs as envs +from vllm.models.qwen4_exp.nvidia import sm70_fp16_gemv as module + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True) +def uncached_environment(): + envs.disable_envs_cache() + + +def metadata(rows, width, **overrides): + values = dict( + ndim=2, + shape=(rows, width), + dtype=torch.float16, + is_cuda=True, + device=torch.device("cuda:0"), + stride=lambda: (width, 1), + ) + return SimpleNamespace(**(values | overrides)) + + +@pytest.mark.parametrize("rows", (0, 1, 2, 3, 4, 8, 16, 17, 32, 64, 127, 256, 4096)) +def test_admission_has_no_max_batch_binding(monkeypatch, rows): + monkeypatch.setenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", "1") + monkeypatch.setattr(module.current_platform, "is_device_capability", lambda _: True) + assert module._can_fuse_gdn_projection_split( + metadata(rows, 4096), metadata(rows, 24) + ) == (rows > 1) + + +@pytest.mark.parametrize( + "change", + ( + dict(dtype=torch.float32), + dict(dtype=torch.bfloat16), + dict(is_cuda=False), + dict(ndim=3), + dict(shape=(4, 4095)), + dict(stride=lambda: (8192, 2)), + dict(device=torch.device("cuda:1")), + ), +) +def test_unsupported_input_falls_back(monkeypatch, change): + monkeypatch.setenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", "1") + monkeypatch.setattr(module.current_platform, "is_device_capability", lambda _: True) + assert not module._can_fuse_gdn_projection_split( + metadata(4, 4096, **change), metadata(4, 24) + ) + + +def test_off_and_non_sm70_fall_back(monkeypatch): + q, ba = metadata(4, 4096), metadata(4, 24) + monkeypatch.setenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", "0") + assert not module._can_fuse_gdn_projection_split(q, ba) + monkeypatch.setenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", "1") + monkeypatch.setattr( + module.current_platform, "is_device_capability", lambda _: False + ) + assert not module._can_fuse_gdn_projection_split(q, ba) + + +def test_default_is_off(monkeypatch): + monkeypatch.delenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", raising=False) + assert not envs.environment_variables["VLLM_SM70_GDN_BATCH_SPLIT_COPY"]() + + +def test_fusion_keeps_both_linear_calls(monkeypatch): + x = torch.empty(4, 2560) + wq, wb = torch.empty(4096, 2560), torch.empty(24, 2560) + q, ba = torch.empty(4, 4096), torch.empty(4, 24) + calls = [] + + def linear(value, weight): + assert value is x + calls.append(weight) + return q if weight is wq else ba + + expected = tuple(torch.empty(4, width) for width in (2560, 1536, 12, 12)) + + def fused(value, gate): + assert value is q and gate is ba + return expected + + monkeypatch.setattr(torch.nn.functional, "linear", linear) + monkeypatch.setattr(module, "_can_fuse_gdn_projection_split", lambda *args: True) + monkeypatch.setattr(module, "_split_gdn_projection_outputs", fused) + actual = module._qwen38_sm70_fp16_gdn_input(x, wq, wb) + assert actual is expected + assert len(calls) == 2 and calls[0] is wq and calls[1] is wb + + +@pytest.mark.parametrize("rows", (0, 1, 4, 17)) +def test_cpu_original_fallback_is_unchanged(rows): + # Small K is sufficient to exercise fallback slicing, including empty input. + x, wq, wb = torch.randn(rows, 3), torch.randn(4096, 3), torch.randn(24, 3) + actual = module._qwen38_sm70_fp16_gdn_input(x, wq, wb) + q, ba = torch.nn.functional.linear(x, wq), torch.nn.functional.linear(x, wb) + expected = q[:, :2560], q[:, 2560:], ba[:, :12], ba[:, 12:] + assert all( + torch.equal(a, b) and a.is_contiguous() + for a, b in zip(actual, expected, strict=True) + ) + + +@pytest.mark.parametrize("rows", (1, 2, 4, 8, 16, 17, 32, 64)) +def test_cuda_public_op_graph_replay_is_bitwise(monkeypatch, rows): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("Requires an owned SM70 GPU") + envs.disable_envs_cache() + torch.manual_seed(51 + rows) + x = torch.randn(rows, 2560, device="cuda", dtype=torch.float16) + wq = torch.randn(4096, 2560, device="cuda", dtype=torch.float16) * 0.01 + wb = torch.randn(24, 2560, device="cuda", dtype=torch.float16) * 0.01 + op = torch.ops.vllm.qwen38_sm70_fp16_gdn_input + route_hits = [] + fused = module._split_gdn_projection_outputs + + def tracked(q, ba): + route_hits.append(q.shape[0]) + return fused(q, ba) + + monkeypatch.setattr(module, "_split_gdn_projection_outputs", tracked) + monkeypatch.setenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", "1") + op(x, wq, wb) + torch.accelerator.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + outputs = op(x, wq, wb) + assert route_hits == ([rows, rows] if rows > 1 else []) + route_hits.clear() + for i in range(4): + x.normal_(0, 0.1 * (i + 1)) + wq.mul_(0.9) + wb.mul_(0.9) + for out in outputs: + out.fill_(float("nan")) + monkeypatch.setenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", "0") + reference = op(x, wq, wb) + g.replay() + assert not route_hits + assert all( + a.is_contiguous() and torch.equal(a.view(torch.int16), b.view(torch.int16)) + for a, b in zip(outputs, reference, strict=True) + ) diff --git a/tools/prepare-flashinfer-sm70-qsa.sh b/tools/prepare-flashinfer-sm70-qsa.sh new file mode 100644 index 0000000000..6b5be5e242 --- /dev/null +++ b/tools/prepare-flashinfer-sm70-qsa.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +set -euo pipefail + +# Separate pin from the older WMMA primitive probe. No wheel/package install. +root=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd) +sha=6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f +cccl=16bd510c9b712e82b0ab6cbb630d8e29ba1f7116 +source_dir=${FLASHINFER_SM70_QSA_SOURCE:-"$root/.deps/flashinfer-6c14bbd5ff34"} +if [[ ! -e "$source_dir" ]]; then + git clone --filter=blob:none --no-checkout \ + https://github.com/flashinfer-ai/flashinfer.git "$source_dir" + git -C "$source_dir" checkout --detach "$sha" +fi +[[ $(git -C "$source_dir" rev-parse HEAD) == "$sha" ]] || { + echo "FlashInfer source has a different revision; refusing to overwrite" >&2 + exit 1 +} +[[ -z $(git -C "$source_dir" status --porcelain --untracked-files=no) ]] || { + echo "FlashInfer source has local edits; refusing to overwrite" >&2 + exit 1 +} +git -C "$source_dir" submodule update --init --depth 1 3rdparty/cccl +[[ $(git -C "$source_dir/3rdparty/cccl" rev-parse HEAD) == "$cccl" ]] +echo "Prepared FlashInfer $sha with CCCL $cccl at $source_dir" diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index a666cc9d34..63d26ca7a4 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3113,6 +3113,26 @@ def supports_sm70_qwen38_hc_output_allgather() -> bool: return hasattr(_custom_ar_owner_namespace(), "sm70_qwen38_hc_output_allgather") +def supports_sm70_qwen38_hc_batch() -> bool: + owner = _custom_ar_owner_namespace() + return all( + hasattr(owner, name) + for name in ("sm70_qwen38_hc_batch_down", "sm70_qwen38_hc_batch_mix") + ) + + +def sm70_qwen38_hc_batch_down( + ptr: int, inp: torch.Tensor, injection: torch.Tensor, lora: torch.Tensor +) -> None: + _custom_ar_owner_namespace().sm70_qwen38_hc_batch_down(ptr, inp, injection, lora) + + +def sm70_qwen38_hc_batch_mix( + ptr: int, gate: torch.Tensor, branches: torch.Tensor, output: torch.Tensor +) -> None: + _custom_ar_owner_namespace().sm70_qwen38_hc_batch_mix(ptr, gate, branches, output) + + def supports_sm70_qwen38_hc_shard() -> bool: owner = _custom_ar_owner_namespace() return all( diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index fc46b2419e..750dd22c3d 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -134,6 +134,7 @@ def __init__( register_nccl_symmetric_ops(self.pynccl_comm) self.ca_comm: CustomAllreduce | None = None + self.sm70_hc_batch_comm: CustomAllreduce | None = None self.qr_comm: QuickAllReduce | None = None self.symm_mem_comm: SymmMemCommunicator | None = None self.fi_ar_comm: FlashInferAllReduce | None = None @@ -176,6 +177,19 @@ def __init__( long_prefill_fusion_enabled=use_sm70_tp4_long_prefill_fused_norm, ) + if ( + envs.VLLM_SM70_QWEN38_BATCH_HC_FP16 + and use_custom_allreduce + and current_platform.is_device_capability(70) + and self.ca_comm.supports_sm70_qwen38_hc_batch() + ): + # One isolated packet/epoch channel per TP group, prepared + # before model loading/capture. Ordinary/M1 collectives retain + # their existing communicator and auxiliary-stream ownership. + self.sm70_hc_batch_comm = CustomAllreduce( + group=self.cpu_group, device=self.device, max_size=128 * 1024 + ) + if current_platform.is_rocm(): # Initialize a custom quick all-reduce implementation for AMD. # Quick reduce is designed as a complement to custom allreduce. @@ -617,6 +631,10 @@ def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor: raise ValueError("No PyNCCL communicator found") def destroy(self): + hc_comm = getattr(self, "sm70_hc_batch_comm", None) + if hc_comm is not None: + hc_comm.close() + self.sm70_hc_batch_comm = None if self.pynccl_comm is not None: self.pynccl_comm.destroy() self.pynccl_comm = None diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 24d4524916..9ce30c3b31 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -467,6 +467,27 @@ def all_reduce_sum2( ops.all_reduce_sum2(self._ptr, inp_a, inp_b, out) return out + def supports_sm70_qwen38_hc_batch(self) -> bool: + return bool( + not self.disabled + and self.world_size == 4 + and self.fully_connected + and self.sm70_tp4_push_buffer_ptrs is not None + and ops.supports_sm70_qwen38_hc_batch() + ) + + def can_sm70_qwen38_hc_batch(self, branches: torch.Tensor) -> bool: + return bool( + self.supports_sm70_qwen38_hc_batch() + and branches.is_cuda + and branches.device == self.device + and branches.dtype == torch.float16 + and branches.ndim == 2 + and 2 <= branches.shape[0] <= 16 + and branches.shape[1] == 10240 + and branches.is_contiguous() + ) + def can_sm70_qwen38_hc_shard(self, branches: torch.Tensor) -> bool: return bool( not self.disabled diff --git a/vllm/envs.py b/vllm/envs.py index f686e9b0e0..3863899895 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -172,7 +172,10 @@ VLLM_SM70_QWEN4_EXP_ONLINE_QPN8: bool = False VLLM_SM70_QWEN38_FP16_GEMV: bool = False VLLM_SM70_QWEN38_FUSED_GDN_INPUT_FP16: bool = False + VLLM_SM70_GDN_BATCH_SPLIT_COPY: bool = False VLLM_SM70_QWEN38_FUSED_HC_FP16: bool = False + VLLM_SM70_QWEN38_BATCH_HC_FP16: bool = False + VLLM_SM70_FLASHINFER_BATCH: bool = False VLLM_SM70_QWEN38_DUAL_COMPILE: bool = False VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION: bool = True VLLM_SM70_FP8_QPN8_PP2_TP4: bool = False @@ -1784,12 +1787,25 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_QWEN38_FUSED_GDN_INPUT_FP16": lambda: bool( int(os.getenv("VLLM_SM70_QWEN38_FUSED_GDN_INPUT_FP16", "0")) ), + # Copy-only batched fallback fusion: keep both GDN GEMMs unchanged. + # Experimental until complete-model performance and quality gates pass. + "VLLM_SM70_GDN_BATCH_SPLIT_COPY": lambda: bool( + int(os.getenv("VLLM_SM70_GDN_BATCH_SPLIT_COPY", "0")) + ), # Fuse the exact Qwen3.8 M=1 HyperConnection down/SiLU and up/gate-mix # stages while retaining FP16 checkpoint weights and inter-stage rounding. # This remains opt-in pending the same model-level quality gates as GEMV. "VLLM_SM70_QWEN38_FUSED_HC_FP16": lambda: bool( int(os.getenv("VLLM_SM70_QWEN38_FUSED_HC_FP16", "0")) ), + # Experimental FP16 HC batch fusion; no model-quality/default admission yet. + "VLLM_SM70_QWEN38_BATCH_HC_FP16": lambda: bool( + int(os.getenv("VLLM_SM70_QWEN38_BATCH_HC_FP16", "0")) + ), + # Integration probe only; keep off until model quality/performance admission. + "VLLM_SM70_FLASHINFER_BATCH": lambda: bool( + int(os.getenv("VLLM_SM70_FLASHINFER_BATCH", "0")) + ), # Exact M=1 Qwen3Next/Qwen4Exp shared-expert output gate. This replaces # the scalar GEMV, sigmoid, and output multiply with one SM70 kernel while # retaining the checkpoint's FP16 accumulation and output rounding. diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 154a85dd03..19353bbed3 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -221,6 +221,43 @@ def is_forward_context_available() -> bool: return _forward_context is not None +def is_uniform_decode_metadata(metadata: Any) -> bool: + """Recognize one-token decode using CPU metadata, without GPU readback. + + List/DBO metadata needs its own per-microbatch ownership. Unknown or + tensor-valued counters fail closed rather than synchronizing the device. + """ + if not isinstance(metadata, dict) or not metadata: + return False + seen_decode = False + for meta in metadata.values(): + prefills = getattr(meta, "num_prefills", 0) + prefill_tokens = getattr(meta, "num_prefill_tokens", 0) + if ( + not isinstance(prefills, int) + or not isinstance(prefill_tokens, int) + or prefills != 0 + or prefill_tokens != 0 + ): + return False + max_query = getattr(meta, "max_query_len", None) + if max_query is not None: + if not isinstance(max_query, int) or max_query != 1: + return False + seen_decode = True + num_decodes = getattr(meta, "num_decodes", None) + if num_decodes is not None: + decode_tokens = getattr(meta, "num_decode_tokens", None) + if ( + not isinstance(num_decodes, int) + or not isinstance(decode_tokens, int) + or decode_tokens != num_decodes + ): + return False + seen_decode |= num_decodes > 0 + return seen_decode + + def create_forward_context( attn_metadata: Any, vllm_config: VllmConfig, diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 685cadbc94..7660f3ca3d 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -4047,7 +4047,12 @@ def forward_cuda( """ num_tokens = hidden_states.size(0) layer_name = _encode_layer_name(self.prefix) - if _sm70_qwen_gdn_input_core_boundary_enabled(): + # Keep the compiled boundary independent of the example prefill size. + # vLLM reuses this graph for decode; try_gdn must inspect the actual + # shape and scheduler metadata INSIDE the opaque op on each capture. + if _sm70_qwen_gdn_input_core_boundary_enabled() or ( + getattr(self, "_sm70_fi_ready", False) and use_sm70_decode_graph_semantics() + ): z = torch.empty( (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), dtype=hidden_states.dtype, @@ -7255,6 +7260,24 @@ def qwen_gdn_input_projection_core( forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] + if getattr(self, "_sm70_fi_ready", False): + from vllm.model_executor.layers.sm70_flashinfer_batch import try_gdn + + raw_metadata = forward_context.attn_metadata + metadata = ( + raw_metadata.get(layer_name) if isinstance(raw_metadata, dict) else None + ) + if try_gdn( + self, + hidden_states, + z_out, + core_attn_out, + conv_state_cache, + ssm_state_cache, + metadata, + ): + return z_out, core_attn_out + hidden_states = _sm70_dump_gdn_projection_tensor( "gdn_hidden_states_input_core", layer_name, @@ -7285,6 +7308,21 @@ def qwen_gdn_input_projection_core( "SM70 GDN QPN8 N4096 plus FP16 b/a N24 split route enabled." ) z = z_out + elif ( + getattr(self, "_sm70_fi_ready", False) + and getattr(self, "sm70_qwen38_fp16_fused_input", False) + and use_sm70_decode_graph_semantics() + and not _sm70_gdn_projection_dump_requested(layer_name) + ): + # The opt-in boundary also sees M1 and prefill. Preserve their existing + # projection implementation, including the fused FP16 M1 kernel, + # instead of silently replacing it with separate QKVZ and BA GEMMs. + mixed_qkv, z, b, a = torch.ops.vllm.qwen38_sm70_fp16_gdn_input( + hidden_states, + self.in_proj_qkvz.weight, + self.in_proj_ba.weight, + ) + z_out.copy_(z.reshape_as(z_out)) else: mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) ba, _ = self.in_proj_ba(hidden_states) diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 85bd286fde..cda20ca943 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -19,7 +19,11 @@ from vllm import _sm70_ops as sm70_ops from vllm import envs from vllm.config.vllm import get_current_vllm_config_or_none -from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.forward_context import ( + get_forward_context, + is_forward_context_available, + is_uniform_decode_metadata, +) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, @@ -272,43 +276,11 @@ def _grouped_decode_context_ok() -> bool: if not is_forward_context_available(): return False context = get_forward_context() - metadata = context.attn_metadata - # DBO/list metadata needs per-microbatch ownership, not a shared decision. - if not isinstance(metadata, dict) or not metadata: - return False key = "sm70_grouped_moe_decode" if key not in context.additional_kwargs: - seen_decode = False - allowed = True - for meta in metadata.values(): - prefills = getattr(meta, "num_prefills", 0) - prefill_tokens = getattr(meta, "num_prefill_tokens", 0) - if ( - not isinstance(prefills, int) - or not isinstance(prefill_tokens, int) - or prefills != 0 - or prefill_tokens != 0 - ): - allowed = False - break - max_query = getattr(meta, "max_query_len", None) - if max_query is not None: - if not isinstance(max_query, int) or max_query != 1: - allowed = False - break - seen_decode = True - num_decodes = getattr(meta, "num_decodes", None) - if num_decodes is not None: - decode_tokens = getattr(meta, "num_decode_tokens", None) - if ( - not isinstance(num_decodes, int) - or not isinstance(decode_tokens, int) - or decode_tokens != num_decodes - ): - allowed = False - break - seen_decode |= num_decodes > 0 - context.additional_kwargs[key] = bool(allowed and seen_decode) + context.additional_kwargs[key] = is_uniform_decode_metadata( + context.attn_metadata + ) return bool(context.additional_kwargs[key]) diff --git a/vllm/model_executor/layers/sm70_flashinfer_batch.py b/vllm/model_executor/layers/sm70_flashinfer_batch.py new file mode 100644 index 0000000000..b2b16a0a9b --- /dev/null +++ b/vllm/model_executor/layers/sm70_flashinfer_batch.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in integration probe for FlashInfer-derived SM70 components. + +Libraries are built and loaded before graph capture. This module does not JIT +compile, quantize weights, change scheduler state ownership, or replace HC/M1. +All scratch tensors are call-local so different graph/ubatch streams cannot +overwrite a shared mutable workspace. Only the zero page and weights persist. +""" + +import importlib.util +import os + +import torch + +from vllm.config import get_current_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.platforms import current_platform + +logger = init_logger(__name__) +_QSA_ZERO: dict[torch.device, torch.Tensor] = {} +_MQA_SMS: dict[torch.device, int] = {} + + +def load_native_fragment(module_name: str, namespaces: tuple[str, ...]) -> None: + # An explicitly preloaded prototype may own one of these namespaces. + # Loading the wheel fragment in that case would register it twice and abort + # the process. Preserve the override and locally fall back on missing shapes. + if any(hasattr(getattr(torch.ops, name), "run") for name in namespaces): + return + spec = importlib.util.find_spec(module_name) + if spec is not None and spec.origin is not None: + torch.ops.load_library(spec.origin) + logger.info("Loaded SM70 FlashInfer native fragment: %s", module_name) + + +def copy_derived_buffer(layer, name, value): + existing = getattr(layer, name, None) + if existing is None: + layer.register_buffer(name, value.clone().contiguous(), persistent=False) + else: + if (existing.shape, existing.dtype, existing.device) != ( + value.shape, + value.dtype, + value.device, + ): + raise RuntimeError( + "FlashInfer derived geometry changed; rebuild CUDA graphs" + ) + existing.copy_(value) + + +def uniform_decode(metadata, rows: int) -> bool: + """Metadata, not a server max-seqs or prefill-budget setting, owns routing.""" + return bool( + metadata is not None + and 2 <= rows <= 64 + and metadata.num_prefills == 0 + and metadata.num_prefill_tokens == 0 + and metadata.num_spec_decodes == 0 + and metadata.num_spec_decode_tokens == 0 + and 0 < metadata.num_decodes <= rows + and metadata.num_decode_tokens == metadata.num_decodes + and metadata.non_spec_state_indices_tensor is not None + and metadata.non_spec_state_indices_tensor.numel() >= rows + ) + + +def prepare(model: torch.nn.Module, device: torch.device) -> None: + if ( + device.type != "cuda" + or not current_platform.is_device_capability(70) + or get_current_vllm_config().speculative_config is not None + ): + return + from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, + ) + + for key in ("VLLM_SM70_FLASHINFER_GDN_LIBRARY", "VLLM_SM70_FLASHINFER_QSA_LIBRARY"): + path = os.environ.get(key) + if path: + torch.ops.load_library(path) + # Wheel-native fragment. Probes may preload the same operator explicitly; + # never load a second definition or compile on a serving worker. + load_native_fragment("vllm._sm70_flashinfer_C", ("_C_flashinfer_mqa_sm70",)) + load_native_fragment( + "vllm._sm70_flashinfer_gdn_C", + tuple( + f"_C_flashinfer_gdn_sm70_h2560_q{q}_v{v}" + for q, v in ((4, 12), (8, 24), (16, 48)) + ), + ) + if hasattr(torch.ops._C_flashinfer_mqa_sm70, "run"): + concrete_device = torch.empty(0, device=device).device + _MQA_SMS[concrete_device] = torch.cuda.get_device_properties( + concrete_device + ).multi_processor_count + if any( + hasattr(getattr(torch.ops, ns), "run") + for ns in ("_C_flashinfer_qsa_sm70_compat", "_C_flashinfer_qsa_sm70") + ): + concrete_device = torch.empty(0, device=device).device + if concrete_device not in _QSA_ZERO: + _QSA_ZERO[concrete_device] = torch.zeros( + 256, dtype=torch.float16, device=concrete_device + ) + count = 0 + for layer in model.modules(): + if not isinstance(layer, QwenGatedDeltaNetAttention): + continue + qh, vh = layer.num_k_heads // layer.tp_size, layer.num_v_heads // layer.tp_size + ba = getattr(layer.in_proj_ba, "weight", None) + qkvz = getattr(layer.in_proj_qkvz, "weight", None) + if ( + ba is None + or qkvz is None + or ba.dtype != torch.float16 + or qkvz.dtype != torch.float16 + or ba.device.type != "cuda" + or qkvz.device != ba.device + or layer.gqa_interleaved_layout + or layer.disable_tp_for_ba_proj + or layer.head_k_dim != 128 + or layer.head_v_dim != 128 + or ba.ndim != 2 + or qkvz.ndim != 2 + or ba.shape != (2 * vh, qkvz.shape[1]) + or qkvz.shape[0] != (2 * qh + 2 * vh) * 128 + or layer.A_log.dtype != torch.float32 + or layer.A_log.numel() != vh + or layer.dt_bias.dtype != torch.float16 + or layer.dt_bias.numel() != vh + or layer.conv1d.weight.dtype != torch.float16 + or layer.conv1d.weight.numel() != (2 * qh + vh) * 128 * 4 + or layer.activation != "silu" + ): + if getattr(layer, "_sm70_fi_ready", False): + raise RuntimeError( + "FlashInfer GDN prepared contract changed; rebuild CUDA graphs" + ) + continue + namespace = f"_C_flashinfer_gdn_sm70_h{ba.shape[1]}_q{qh}_v{vh}" + ops = getattr(torch.ops, namespace) + if not hasattr(ops, "run"): + if getattr(layer, "_sm70_fi_ready", False): + raise RuntimeError( + "FlashInfer GDN prepared geometry unavailable; rebuild CUDA graphs" + ) + continue + copy_derived_buffer(layer, "_sm70_fi_ba", ba.t()) + copy_derived_buffer( + layer, + "_sm70_fi_bias", + layer.conv1d.bias if layer.conv1d.bias is not None else ba.new_empty(0), + ) + layer._sm70_fi_op = ops.run + layer._sm70_fi_ready = True + count += 1 + logger.info( + "FlashInfer SM70 experimental capabilities: GDN layers=%d, MQA=%s, " + "sparse QSA=%s; missing components use " + "local fallback, no M1/HC change.", + count, + hasattr(torch.ops._C_flashinfer_mqa_sm70, "run"), + any( + hasattr(getattr(torch.ops, ns), "run") + for ns in ("_C_flashinfer_qsa_sm70_compat", "_C_flashinfer_qsa_sm70") + ), + ) + + +def try_gdn(layer, hidden, z_out, core_out, conv_cache, state, metadata) -> bool: + rows = hidden.shape[0] + if not getattr(layer, "_sm70_fi_ready", False) or not uniform_decode( + metadata, rows + ): + return False + # As in the existing recurrent core, AOT tracing before KV allocation + # can pass empty placeholders. Resolve only those placeholders from the + # scheduler-bound layer inside the opaque runtime boundary. Never replace + # explicit state arguments (e.g. a graph/ubatch's dedicated cache). + if conv_cache.numel() == 0 and state.numel() == 0: + cache = getattr(layer, "kv_cache", None) + if cache is not None and cache[0].numel() > 0: + conv_cache, state = cache[0], cache[1] + if conv_cache.ndim != 3 or state.ndim != 4: + return False + conv = conv_cache if is_conv_state_dim_first() else conv_cache.transpose(-1, -2) + vh = layer.num_v_heads // layer.tp_size + qh = layer.num_k_heads // layer.tp_size + width = (2 * qh + vh) * 128 + if ( + hidden.dtype != torch.float16 + or not hidden.is_contiguous() + or state.dtype != torch.float32 + or state.ndim != 4 + or state.shape[1:] != (vh, 128, 128) + or state.stride()[1:] != (128 * 128, 128, 1) + or conv.dtype != torch.float16 + or conv.ndim != 3 + or conv.shape[1:] != (width, 3) + or conv.shape[0] != state.shape[0] + ): + return False + # This boundary owns both projection and recurrence: do not also compute + # the separate BA GEMM, which would erase the gate-fusion benefit. + qkvz, _ = layer.in_proj_qkvz(hidden) + qkv = qkvz[:, :width] + z_out.copy_(qkvz[:, width:].reshape_as(z_out)) + conv_weight = layer.conv1d.weight.reshape(width, 4) + conv_out = hidden.new_empty((rows, width)) + partial = torch.empty( + rows * 2 * vh * 160, device=hidden.device, dtype=torch.float32 + ) + layer._sm70_fi_op( + hidden, + layer._sm70_fi_ba, + qkv, + conv_weight, + layer._sm70_fi_bias, + conv, + layer.A_log, + layer.dt_bias, + state, + metadata.non_spec_state_indices_tensor[:rows], + core_out, + conv_out, + partial, + ) + logger.info_once("Selected FlashInfer SM70 fused GDN batch route, rows=%d.", rows) + return True + + +def try_qsa(q, k, v, indices, table, requests, out): + zero = _QSA_ZERO.get(q.device) + if ( + zero is None + or not 4 <= q.shape[0] <= 16 + or q.dtype != torch.float16 + or k.dtype != torch.float16 + or v.dtype != torch.float16 + or q.shape[2] != 256 + or q.shape[1] > 32 + or q.shape[1] // k.shape[2] not in (1, 2, 4, 6, 8) + or not out.is_contiguous() + or k.stride() != v.stride() + or any( + t.data_ptr() % 16 or any(s % 8 for s in t.stride()[:-1]) for t in (q, k, v) + ) + ): + return None + rows, heads, dim = q.shape + compatible = hasattr(torch.ops._C_flashinfer_qsa_sm70_compat, "run") + if compatible: + from vllm.models.qwen4_exp.nvidia.ops.qsa import _qsa_sparse_launch_profile + + # Match the production 16-token tile partition as well as FP16 P. + # The SIMT experiment uses a different split profile and remains + # available only when explicitly preloaded for counterfactual audits. + group = heads // k.shape[2] + block_n, target_splits, _ = _qsa_sparse_launch_profile( + rows * k.shape[2], 1 << (group - 1).bit_length(), True + ) + if block_n != 16 or indices.shape[1] <= 0: + return None + tiles = (indices.shape[1] + block_n - 1) // block_n + splits = min(target_splits, 1 << (tiles.bit_length() - 1)) + op = torch.ops._C_flashinfer_qsa_sm70_compat.run + else: + splits = 16 if rows >= 16 else 32 + op = torch.ops._C_flashinfer_qsa_sm70.run + width = ((indices.shape[1] + splits - 1) // splits) * splits + offsets = torch.empty((rows, width), device=q.device, dtype=torch.int64) + metadata = torch.empty( + rows + 2 + 2 * rows * splits, device=q.device, dtype=torch.int32 + ) + partial = torch.empty( + (rows, splits, heads, dim), device=q.device, dtype=torch.float32 + ) + lse = torch.empty((rows, splits, heads), device=q.device, dtype=torch.float32) + op( + q, + k, + v, + indices, + table, + requests, + offsets, + metadata, + zero, + partial, + lse, + out, + splits, + ) + logger.info_once( + "Selected FlashInfer SM70 sparse QSA batch route, rows=%d, " + "FP16-P compatibility=%s.", + rows, + compatible, + ) + return out + + +def try_mqa(q, k, table, requests, positions, lengths, ratio, divisor, out, visible): + """Opt-in device-planned scorer, called inside the opaque QSA boundary.""" + sms = _MQA_SMS.get(q.device) + if ( + sms is None + or not 4 <= q.shape[0] <= 16 + or q.dtype != torch.float16 + or k.dtype != torch.float16 + or q.shape[2] != 128 + or q.shape[1] != 4 + or q.stride(2) != 1 + or k.stride(3) != 1 + or any(t.device != q.device for t in (k, table, requests, positions, lengths)) + or any(t.dtype != torch.int32 for t in (table, requests, lengths)) + or positions.dtype not in (torch.int32, torch.int64) + or any(not t.is_contiguous() for t in (requests, positions, lengths)) + or table.stride(1) != 1 + or any(t.data_ptr() % 16 or any(s % 8 for s in t.stride()[:-1]) for t in (q, k)) + ): + return False + workers = sms * (4 if q.shape[0] >= 16 else 2) + schedule = torch.empty((workers + 1, 2), dtype=torch.int32, device=q.device) + torch.ops._C_flashinfer_mqa_sm70.run( + q, + k, + table, + requests, + positions, + lengths, + out, + visible, + schedule, + ratio, + divisor, + workers, + ) + logger.info_once( + "Selected FlashInfer SM70 device-planned QSA MQA scorer, rows=%d.", + q.shape[0], + ) + return True diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 3b4fcfa042..79495a22d0 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -121,6 +121,15 @@ def process_weights_after_loading( with device_loading_context(module, target_device): module.process_weights_after_loading(model_config.dtype) + # Model-owned derived weights must observe the final quantization layout. + # This optional SM70 hook is a no-op unless the experimental route is enabled. + if prepare_batch_hc := getattr(model, "prepare_sm70_batch_hc", None): + prepare_batch_hc() + if envs.VLLM_SM70_FLASHINFER_BATCH: + from vllm.model_executor.layers.sm70_flashinfer_batch import prepare + + prepare(model, target_device) + # Needed for torchao model reloading via model.reload_weights # @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights` if model_config.quantization == "torchao": diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py index 787865f687..b8b0f88c3e 100644 --- a/vllm/models/qwen4_exp/nvidia/hyperconnection.py +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -45,6 +45,7 @@ hc_gate_mix, hc_silu, ) +from .sm70_batch_hc import maybe_apply_sm70_batch_hc from .sm70_fp16_hc import maybe_apply_qwen38_sm70_fp16_fused_hc @@ -130,6 +131,9 @@ def __init__( def _project(self, xn: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: if self.use_combine: + batch_hc = maybe_apply_sm70_batch_hc(self, xn) + if batch_hc is not None: + return batch_hc fused_fp16 = maybe_apply_qwen38_sm70_fp16_fused_hc( self.input_mix_weight_down_block_inject, self.input_mix_weight_up, diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index e88242d661..9cdec4caa6 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -840,6 +840,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) object.__setattr__(self, "_sm70_decode_graph_model", None) + def prepare_sm70_batch_hc(self) -> None: + from .sm70_batch_hc import prepare_sm70_batch_hc + + prepare_sm70_batch_hc(self) + def prepare_sm70_decode_graph_model(self) -> bool: """Create the shared-weight decode compiler just before graph capture.""" if not envs.VLLM_SM70_QWEN38_DUAL_COMPILE: @@ -1208,6 +1213,9 @@ def prepare_sm70_decode_graph_model(self) -> bool: """Forward decode compiler setup to the wrapped language model.""" return self.language_model.prepare_sm70_decode_graph_model() + def prepare_sm70_batch_hc(self) -> None: + self.language_model.prepare_sm70_batch_hc() + def forward( self, input_ids: torch.Tensor | None, diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index c02b3505c3..0d929475af 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -10,6 +10,7 @@ import regex as re import torch +import vllm.envs as envs from vllm.logger import init_logger from vllm.models.deepseek_v4.common.ops.fp8_software import ( fp8_e4m3fn_bits_to_fp32_bitcast as fp8_e4m3fn_bits_to_fp32, @@ -443,7 +444,6 @@ def _qsa_xqa_page4_table_kernel( OUTPUT_PAGES: tl.constexpr, BLOCK_PAGES: tl.constexpr, PHYSICAL_PAGE_STRIDE: tl.constexpr, - TAIL_MARKER: tl.constexpr, ) -> None: row = tl.program_id(0) slots = tl.arange(0, BLOCK_PAGES) @@ -508,13 +508,16 @@ def _qsa_xqa_page4_table_kernel( physical_microblock = ( tl.maximum(physical_page, 0) * PHYSICAL_PAGE_STRIDE + page_offset // 4 ) + # Sort by logical token, not allocator-dependent physical page ID. Keep + # the partial causal page after all complete pages and invalid slots last. + logical_key = safe_token.to(tl.int64) << 31 encoded = tl.where( valid & is_complete, - physical_microblock, + logical_key | physical_microblock.to(tl.int64), tl.where( valid & is_tail, - physical_microblock + TAIL_MARKER, - 2147483647, + (1 << 62) | logical_key | physical_microblock.to(tl.int64), + 9223372036854775807, ), ) tl.store( @@ -1069,6 +1072,22 @@ def qsa_mqa_paged( visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) if not q.shape[0] or not columns: return logits, visible_blocks + if envs.VLLM_SM70_FLASHINFER_BATCH: + from vllm.model_executor.layers.sm70_flashinfer_batch import try_mqa + + if try_mqa( + q, + k_cache, + page_table, + token_to_req, + query_positions, + sequence_lengths, + compress_ratio, + float(score_divisor), + logits, + visible_blocks, + ): + return logits, visible_blocks sm70_single_token = q.shape[0] == 1 and current_platform.is_device_capability(70) # On V100 the GB300 decode tile leaves the 128-d scorer badly # under-occupied. A 32-column, two-warp tile preserves the selected QSA @@ -1561,7 +1580,7 @@ def _qsa_xqa_page4_block_table( rows = logical_indices.shape[0] encoded_pages = torch.empty( (rows, _SM70_QSA_XQA_PAGE4_PAGES), - dtype=torch.int32, + dtype=torch.int64, device=logical_indices.device, ) xqa_sequence_lengths = torch.empty( @@ -1587,14 +1606,13 @@ def _qsa_xqa_page4_block_table( OUTPUT_PAGES=_SM70_QSA_XQA_PAGE4_PAGES, BLOCK_PAGES=1024, PHYSICAL_PAGE_STRIDE=physical_page_stride, - TAIL_MARKER=_SM70_QSA_XQA_PAGE4_MARKER, num_warps=4, ) sorted_pages = torch.sort(encoded_pages, dim=1).values physical_pages = torch.bitwise_and( sorted_pages, _SM70_QSA_XQA_PAGE4_MARKER - 1, - ) + ).to(torch.int32) return physical_pages, xqa_sequence_lengths @@ -2077,6 +2095,18 @@ def qsa_sparse_paged_attention( if not q.shape[0]: return out + if envs.VLLM_SM70_FLASHINFER_BATCH and not kv_e4m3: + from vllm.model_executor.layers.sm70_flashinfer_batch import try_qsa + + fi_output = try_qsa( + q, k_cache, v_cache, logical_indices, block_table, token_to_req, out + ) + if fi_output is not None: + # Retain the existing FP16 materialization before the FP32 gate. + if output_gate_view is not None: + _qsa_output_gate(fi_output, output_gate_view) + return fi_output + if _use_sm70_qsa_xqa_page4( q, k_cache, diff --git a/vllm/models/qwen4_exp/nvidia/sm70_batch_hc.py b/vllm/models/qwen4_exp/nvidia/sm70_batch_hc.py new file mode 100644 index 0000000000..670e85346b --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/sm70_batch_hc.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in FP16 batched HC with an isolated TP communication channel.""" + +import torch +from torch import nn + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.compilation.sm70_decode_graph import use_sm70_decode_graph_semantics +from vllm.forward_context import ( + get_forward_context, + is_forward_context_available, + is_uniform_decode_metadata, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, + UnquantizedLinearMethod, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op + +from .ops.hc import hc_gate_mix, hc_silu +from .sm70_fp16_gemv import Qwen38SM70FP16LinearMethod, _qwen38_sm70_fp16_gemv +from .sm70_fp16_hc import _qwen38_sm70_fp16_fused_hc + +logger = init_logger(__name__) + + +def _channel(): + from vllm.distributed.parallel_state import get_tp_group + + try: + device_comm = get_tp_group().device_communicator + return getattr(device_comm, "sm70_hc_batch_comm", None) + except (AssertionError, AttributeError, RuntimeError): + return None + + +def _decode_context_ok() -> bool: + if not is_forward_context_available(): + return False + context = get_forward_context() + key = "sm70_batch_hc_decode" + if key not in context.additional_kwargs: + context.additional_kwargs[key] = is_uniform_decode_metadata( + context.attn_metadata + ) + return bool(context.additional_kwargs[key]) + + +def _supported_layers(child: nn.Module) -> bool: + down = getattr(child, "input_mix_weight_down_block_inject", None) + up = getattr(child, "input_mix_weight_up", None) + # Unknown methods or LoRA wrappers must retain their own forward semantics. + return ( + getattr(child, "use_combine", False) + and type(down) is MergedColumnParallelLinear + and type(up) is ReplicatedLinear + and type(down.quant_method) + in (UnquantizedLinearMethod, Qwen38SM70FP16LinearMethod) + and type(up.quant_method) is UnquantizedLinearMethod + and ( + getattr(child, "_sm70_qwen38_fp16_fused_hc", False) + or not any( + getattr(layer, "_sm70_f16_prepared", False) for layer in (down, up) + ) + ) + ) + + +def _supported_weights(down: torch.Tensor, up: torch.Tensor) -> bool: + return ( + down.shape == (336, 10240) + and up.shape == (10240, 320) + and down.dtype == up.dtype == torch.float16 + and down.device == up.device + and down.is_contiguous() + and up.is_contiguous() + ) + + +def _copy_up_shard(child: nn.Module, up: torch.Tensor, rank: int) -> None: + shard = up.view(4, 2560, 320)[:, rank * 640 : (rank + 1) * 640] + packed = shard.reshape(2560, 320).contiguous() + existing = getattr(child, "_sm70_batch_hc_up", None) + if existing is not None: + if ( + existing.shape != packed.shape + or existing.dtype != packed.dtype + or existing.device != packed.device + ): + raise RuntimeError( + "HC reload changed a captured shard layout; rebuild graphs" + ) + existing.copy_(packed) + else: + child.register_buffer("_sm70_batch_hc_up", packed, persistent=False) + + +@torch.no_grad() +def prepare_sm70_batch_hc(module: nn.Module) -> None: + """Run after all quantization post-load hooks, never during a forward.""" + if not envs.VLLM_SM70_QWEN38_BATCH_HC_FP16: + return + if not current_platform.is_device_capability(70): + return + channel = _channel() + if channel is None or not channel.supports_sm70_qwen38_hc_batch(): + logger.warning_once("SM70 batch HC unavailable: retaining original HC path.") + return + count = 0 + for child in module.modules(): + if not _supported_layers(child): + continue + down = child.input_mix_weight_down_block_inject.weight + up = child.input_mix_weight_up.weight + if not (down.is_cuda and _supported_weights(down, up)): + continue + if any(getattr(w, "_vllm_is_uva_offloaded", False) for w in (down, up)): + continue + _copy_up_shard(child, up, channel.rank) + count += 1 + logger.info_once( + "Prepared %d experimental SM70 FP16 batch HC up shards (%.1f MiB/rank); " + "M1/prefill fallbacks preserved, dedicated communication channel.", + count, + count * 2560 * 320 * 2 / 1024**2, + ) + + +def _batch_hc( + x: torch.Tensor, + down: torch.Tensor, + up: torch.Tensor, + packed_up: torch.Tensor, + legacy_fused: bool, + down_gemv: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + # Keep all token-count decisions inside the opaque op. A prefill-first + # dynamic compile must not bake its fallback into later decode graphs. + channel = None + if ( + not envs.VLLM_BATCH_INVARIANT + and x.ndim == 2 + and 2 <= x.shape[0] <= 16 + and _supported_weights(down, up) + and x.device == down.device == packed_up.device + and packed_up.shape == (2560, 320) + and packed_up.dtype == torch.float16 + and packed_up.is_contiguous() + and _decode_context_ok() + ): + channel = _channel() + if channel is not None and channel.can_sm70_qwen38_hc_batch(x): + # Preserve the original down GEMM geometry and FP16 rounding. Sharding + # its output changes the cuBLAS reduction and compounds across HC + # modules. Only shard the up projection, whose input is now identical + # to the original path on every rank. + projected = torch.nn.functional.linear(x, down) + lora = hc_silu(projected[:, :320], 4) + injection = projected[:, 320:324].contiguous() + block = x.new_empty((x.shape[0], 2560)) + gate = torch.nn.functional.linear(lora, packed_up) + ops.sm70_qwen38_hc_batch_mix(channel._ptr, gate, x, block) + logger.info_once( + "Using experimental SM70 batch HC with replicated down and sharded up " + "(rows=%d).", + x.shape[0], + ) + return block, injection + + if legacy_fused: + return _qwen38_sm70_fp16_fused_hc(x, down, up) + projected = ( + _qwen38_sm70_fp16_gemv(x, down) + if down_gemv + else torch.nn.functional.linear(x, down) + ) + lora = hc_silu(projected[:, :320], 4) + return hc_gate_mix(x, torch.nn.functional.linear(lora, up), 4), projected[ + :, 320:324 + ] + + +def _batch_hc_fake( + x: torch.Tensor, + down: torch.Tensor, + up: torch.Tensor, + packed_up: torch.Tensor, + legacy_fused: bool, + down_gemv: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + return x.new_empty((x.shape[0], 2560)), x.new_empty((x.shape[0], 4)) + + +direct_register_custom_op( + op_name="qwen38_sm70_batch_hc", + op_func=_batch_hc, + fake_impl=_batch_hc_fake, +) + + +def maybe_apply_sm70_batch_hc(child: nn.Module, x: torch.Tensor): + packed = getattr(child, "_sm70_batch_hc_up", None) + if ( + packed is None + or envs.VLLM_BATCH_INVARIANT + or not use_sm70_decode_graph_semantics() + ): + return None + if not _supported_layers(child): + return None + down_layer = child.input_mix_weight_down_block_inject + if not _supported_weights(down_layer.weight, child.input_mix_weight_up.weight): + return None + return torch.ops.vllm.qwen38_sm70_batch_hc( + x, + down_layer.weight, + child.input_mix_weight_up.weight, + packed, + getattr(child, "_sm70_qwen38_fp16_fused_hc", False), + type(down_layer.quant_method) is Qwen38SM70FP16LinearMethod, + ) diff --git a/vllm/models/qwen4_exp/nvidia/sm70_fp16_gemv.py b/vllm/models/qwen4_exp/nvidia/sm70_fp16_gemv.py index 5fe5182f05..23df02f8ab 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_fp16_gemv.py +++ b/vllm/models/qwen4_exp/nvidia/sm70_fp16_gemv.py @@ -58,6 +58,69 @@ class _GemvPlan(NamedTuple): _SHAPE_PLANS = {shape: plan for _, shape, plan in _ROLE_PLANS} +@triton.jit +def _qwen38_gdn_projection_split_kernel( + qkvz, + ba, + qkv, + z, + b, + a, + QKV: tl.constexpr, + Z: tl.constexpr, + B: tl.constexpr, + A: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + value = tl.load(qkvz + row * (QKV + Z) + col, col < QKV + Z, other=0) + tl.store(qkv + row * QKV + col, value, col < QKV) + tl.store(z + row * Z + col - QKV, value, (col >= QKV) & (col < QKV + Z)) + if tl.program_id(1) == 0: + tl.static_assert(B + A <= BLOCK) + gate_col = tl.arange(0, BLOCK) + gate = tl.load(ba + row * (B + A) + gate_col, gate_col < B + A, other=0) + tl.store(b + row * B + gate_col, gate, gate_col < B) + tl.store(a + row * A + gate_col - B, gate, (gate_col >= B) & (gate_col < B + A)) + + +def _split_gdn_projection_outputs(qkvz, ba): + m = qkvz.shape[0] + out = tuple(qkvz.new_empty((m, n)) for n in (2560, 1536, 12, 12)) + _qwen38_gdn_projection_split_kernel[(m, triton.cdiv(4096, 256))]( + qkvz, + ba, + *out, + QKV=2560, + Z=1536, + B=12, + A=12, + BLOCK=256, + num_warps=4, + num_stages=1, + ) + return out + + +def _can_fuse_gdn_projection_split(qkvz: torch.Tensor, ba: torch.Tensor) -> bool: + # This copy-only operation does not change either GEMM. Keep the existing + # M1 path and all unsupported layouts; no maximum batch/sequence binding. + return bool( + envs.VLLM_SM70_GDN_BATCH_SPLIT_COPY + and _is_packed_row_major(qkvz) + and _is_packed_row_major(ba) + and qkvz.shape[0] > 1 + and qkvz.shape[1] == 4096 + and ba.shape == (qkvz.shape[0], 24) + and qkvz.dtype == ba.dtype == torch.float16 + and qkvz.is_cuda + and ba.is_cuda + and qkvz.device == ba.device + and current_platform.is_device_capability(70) + ) + + @triton.jit def _qwen38_fp16_row_gemv_kernel( x_ptr, @@ -204,6 +267,9 @@ def _qwen38_sm70_fp16_gdn_input( ): qkvz = torch.nn.functional.linear(x, qkvz_weight) ba = torch.nn.functional.linear(x, ba_weight) + if _can_fuse_gdn_projection_split(qkvz, ba): + logger.info_once("SM70 GDN batched projection split-copy fusion enabled.") + return _split_gdn_projection_outputs(qkvz, ba) return ( qkvz[..., :2560].contiguous(), qkvz[..., 2560:].contiguous(), diff --git a/vllm/third_party/flashinfer_sm70/NOTICE.txt b/vllm/third_party/flashinfer_sm70/NOTICE.txt new file mode 100644 index 0000000000..f21ffcd729 --- /dev/null +++ b/vllm/third_party/flashinfer_sm70/NOTICE.txt @@ -0,0 +1,23 @@ +FlashInfer SM70 adapters in 1Cat-vLLM + +Distributed under the Apache License, Version 2.0 (see the distribution LICENSE). + +The paged MQA scheduling algorithm is adapted from FlashInfer, commit +6c14bbd5ff34210404d5d4b5f6ff3b4b2527f59f: + flashinfer/attn_scores/kernels/schedule_kernel.py +Copyright (c) 2025 FlashInfer team. +That source is adapted from TensorRT-LLM's PagedMQALogitsMetadataKernel. +Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +The native FP16 Volta QSA scorer, runtime integration and modifications are +Copyright contributors to the vLLM project, distributed under Apache-2.0. + +The fused SM70 GDN implementation is adapted from FlashInfer at the same +commit, flashinfer/gdn_kernels/experimental/kernel/gdn_fused_decode_sm120.cu +(including the surrounding gate and convolution fusion). +Copyright (c) 2026 FlashInfer team, Apache-2.0. +SM70 cooperative synchronization, FP16 materialization, FP32 state handling, +geometry instantiations and runtime adapters are 1Cat-vLLM modifications. + +This is an architecture-specific adaptation maintained by 1Cat-vLLM, not +an assertion of official FlashInfer SM70 support or upstream endorsement.