diff --git a/benchmarks/kernels/benchmark_sm70_hc_full_chain.py b/benchmarks/kernels/benchmark_sm70_hc_full_chain.py new file mode 100644 index 0000000000..2b29c0f784 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_hc_full_chain.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Complete Qwen3.8 HC workload, including norms and the final mixer. + +Attention/MoE outputs are fixed external inputs; their computation and PLE +are excluded. This is a CUDA Graph microbenchmark, NOT full-model TPOT or +full-model Nsight service time. Run on four exclusively available SM70 GPUs: + +CUDA_VISIBLE_DEVICES=0,1,2,3 CUDA_DEVICE_ORDER=PCI_BUS_ID \ +VLLM_SM70_TP4_PUSH_ALLREDUCE=1 \ +.venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=4 \ + benchmarks/kernels/benchmark_sm70_hc_full_chain.py --model MODEL --out RESULT +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +from pathlib import Path +from statistics import median +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.distributed as dist +from safetensors import safe_open + +import vllm.envs as envs +from benchmarks.kernels.benchmark_sm70_hc_tp4 import load_weights +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce +from vllm.models.qwen4_exp.nvidia.ops.hc import ( + grouped_gemma_rmsnorm, + hc_combine, + hc_combine_norm, + hc_gate_mix, + hc_silu, +) + + +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("--quality-inputs", type=int, default=16) + parser.add_argument("--warmup", type=int, default=1000) + parser.add_argument("--replays", type=int, default=150) + parser.add_argument( + "--fused-up", + action="store_true", + help="Compare hidden split against fused up/mix/gather", + ) + parser.add_argument( + "--aux-stress-replays", + type=int, + default=32, + help="Auxiliary sum2 replays per changing input with --fused-up", + ) + args = parser.parse_args() + 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" + ) + visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if len(visible.split(",")) != 4 or int(os.environ["WORLD_SIZE"]) != 4: + raise RuntimeError("Set CUDA_VISIBLE_DEVICES and launch exactly four ranks") + rank, local_rank = int(os.environ["RANK"]), int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + if torch.cuda.get_device_capability() != (7, 0): + raise RuntimeError("This gate is specific to SM70") + dist.init_process_group("nccl") + group = dist.new_group(backend="gloo") + comm = CustomAllreduce(group=group, device=local_rank, max_size=8 * 1024 * 1024) + try: + owned_pids = [None] * 4 + dist.all_gather_object(owned_pids, os.getpid(), group=group) + + def ensure_exclusive() -> None: + contenders = [] + if rank == 0: + probe = subprocess.check_output( + [ + "nvidia-smi", + "-i", + visible, + "--query-compute-apps=pid,used_memory", + "--format=csv,noheader,nounits", + ], + text=True, + ) + for line in probe.splitlines(): + pid, memory = (int(part.strip()) for part in line.split(",")) + if pid not in owned_pids and memory > 128: + contenders.append({"pid": pid, "MiB": memory}) + shared = [contenders] + dist.broadcast_object_list(shared, src=0, group=group) + if shared[0]: + raise RuntimeError(f"GPU contention invalidates timing: {shared[0]}") + + ensure_exclusive() + if not comm.supports_sm70_qwen38_hc_output_allgather(): + raise RuntimeError("Load a source-matched custom-AR extension") + if args.fused_up and not comm.supports_sm70_qwen38_hc_up_mix_allgather(): + raise RuntimeError("Load an extension with fused HC up/mix/gather") + weights = load_weights(args.model) + mapping = json.loads((args.model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + + def get(name: str) -> torch.Tensor: + with safe_open( + args.model / mapping[name], framework="pt", device="cpu" + ) as f: + return f.get_tensor(name).half().cuda() + + prefix = "model.language_model." + norms = [ + get(f"{prefix}layers.{layer}.{role}_hyper_connection.hc_norm.weight") + for layer in range(48) + for role in ("attn", "mlp") + ] + final_norm = get(prefix + "hyper_connection_mixer.hc_norm.weight") + final_down = get(prefix + "hyper_connection_mixer.input_mix_weight_down.weight") + 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 + ) + cores = torch.randn( + (96, 1, 2560), device="cuda", dtype=torch.float16, generator=gen + ) + if 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 args.fused_up: + sum_gen = torch.Generator(device="cuda").manual_seed(20260905 + rank) + sum_a = torch.randn( + 96, 2560, device="cuda", dtype=torch.float16, generator=sum_gen + ) + sum_b = torch.randn( + 96, 2560, device="cuda", dtype=torch.float16, generator=sum_gen + ) + peer_sums = [torch.empty_like(sum_a) for _ in range(4)] + dist.all_gather(peer_sums, sum_a + sum_b) + expected_sum = torch.zeros_like(sum_a, dtype=torch.float32) + for peer in peer_sums: + expected_sum.add_(peer.float()) + expected_sum = expected_sum.half() + + def finish(state: torch.Tensor, injection: torch.Tensor): + combined, xn = hc_combine_norm( + state, cores[-1], injection, final_norm, 1e-6, 4 + ) + lora = hc_silu(torch.nn.functional.linear(xn, final_down), 4) + gate = torch.nn.functional.linear(lora, final_up) + 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)) + torch.cuda.synchronize() + + def capture(mode: str, overlap: bool = False): + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + outputs = [] + sums = [] + aux = torch.cuda.Stream() if overlap else None + with ( + patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp), + patch.object( + comm, + "supports_sm70_qwen38_hc_output_allgather", + return_value=mode != "gate", + ), + patch.object( + comm, + "supports_sm70_qwen38_hc_up_mix_allgather", + return_value=mode == "fused", + ), + comm.capture(), + torch.cuda.graph(graph), + ): + main_stream = torch.cuda.current_stream() + if aux is not None: + aux.wait_stream(main_stream) + state, injection = initial, None + for i, (down, up) in enumerate(weights): + # PLE at decoder layer 2 requires a materialized state. + # Exclude PLE computation, but retain its HC boundary. + 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 = 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): + sums.append(comm.all_reduce_sum2(sum_a[i], sum_b[i])) + outputs.extend(finish(state, injection)) + if aux is not None: + main_stream.wait_stream(aux) + torch.cuda.synchronize() + dist.barrier() + return graph, outputs, sums + + timed_modes = ("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) + mismatches = {mode: 0 for mode in list(graphs)[1:]} + sum_mismatches = 0 + + def replay_and_check(stress: bool): + for mode, (graph, _, _) in graphs.items(): + repeats = ( + args.aux_stress_replays if stress and mode == "fused_aux" else 1 + ) + for _ in range(repeats): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + expected = torch.cat( + [x.flatten().view(torch.int16) for x in graphs[timed_modes[0]][1]] + ) + diffs = {} + for mode in list(graphs)[1:]: + actual = torch.cat( + [x.flatten().view(torch.int16) for x in graphs[mode][1]] + ) + diffs[mode] = int(torch.count_nonzero(expected != actual)) + sum_diff = 0 + if args.fused_up: + actual_sum = torch.stack(graphs["fused_aux"][2]) + sum_diff = int( + torch.count_nonzero( + actual_sum.view(torch.int16) != expected_sum.view(torch.int16) + ) + ) + return diffs, sum_diff + + for case in range(args.quality_inputs): + initial.normal_(generator=gen) + cores.normal_(generator=gen) + if case == 0: + initial.zero_() + cores.zero_() + elif case == 1: + initial.mul_(0.01) + cores.mul_(0.01) + diffs, sum_diff = replay_and_check(stress=True) + for mode, diff in diffs.items(): + mismatches[mode] += diff + sum_mismatches += sum_diff + quality = [None] * 4 + dist.all_gather_object( + quality, + { + "rank": rank, + "mismatches": sum(mismatches.values()) + sum_mismatches, + "hc_mismatches": mismatches, + "sum2_mismatches": sum_mismatches, + }, + group=group, + ) + if rank == 0: + print({"quality": quality}, flush=True) + if any(q["mismatches"] for q in quality): + raise RuntimeError("Full HC outputs are not bitwise") + ensure_exclusive() + for mode in timed_modes: + graph = graphs[mode][0] + for _ in range(args.warmup): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + samples = {mode: [] for mode in timed_modes} + for repeat in range(3): + modes = timed_modes if repeat % 2 == 0 else timed_modes[::-1] + for mode in modes: + ensure_exclusive() + graph = graphs[mode][0] + for _ in range(20): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + for _ in range(args.replays): + graph.replay() + end.record() + end.synchronize() + times = [None] * 4 + dist.all_gather_object( + times, start.elapsed_time(end) / args.replays, group=group + ) + samples[mode].append(max(times)) + ensure_exclusive() + diffs, sum_diff = replay_and_check(stress=False) + post_quality = [None] * 4 + dist.all_gather_object( + post_quality, + {"rank": rank, "hc_mismatches": diffs, "sum2_mismatches": sum_diff}, + group=group, + ) + if any( + any(q["hc_mismatches"].values()) or q["sum2_mismatches"] + for q in post_quality + ): + raise RuntimeError("Post-timing HC/sum2 output differs after epoch wrap") + if rank == 0: + result = { + "source_sha": subprocess.check_output( + ["git", "rev-parse", "HEAD"], text=True + ).strip(), + "scope": ( + "full semantic HC microbenchmark; " + "excludes attention/MoE/PLE computation" + ), + "counts": { + "mix_pairs": 96, + "combine_norm": 95, + "separate_combine": 1, + "grouped_norm": 2, + "final_projection_pairs": 1, + "final_gate_mix": 1, + }, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "visible_devices": visible, + "quality": quality, + "quality_inputs": args.quality_inputs, + "post_timing_quality": post_quality, + "aux_stress_replays": args.quality_inputs * args.aux_stress_replays + if args.fused_up + else 0, + "samples_ms": samples, + "median_ms": {mode: median(values) for mode, values in samples.items()}, + } + args.out.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2), flush=True) + finally: + comm.close() + dist.destroy_process_group(group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_hc_tp4.py b/benchmarks/kernels/benchmark_sm70_hc_tp4.py new file mode 100644 index 0000000000..a08b08bec8 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_hc_tp4.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exact TP4 HC Mix gate using all 96 Qwen3.8 checkpoint weight pairs. + +Run with torchrun --standalone --nproc-per-node=4 and --model /path/to/model. +Requires four peer-connected SM70 GPUs, VLLM_SM70_TP4_PUSH_ALLREDUCE=1, +and a source-matched custom-AR extension. Does not load the whole model. +Reports Mix-only CUDA Graph time, NOT full HC, TPOT, or service throughput. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +from statistics import median +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.distributed as dist +from safetensors import safe_open + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce +from vllm.models.qwen4_exp.nvidia import sm70_fp16_hc # noqa: F401 + +MODULES = 96 + + +def load_weights(model: Path) -> list[tuple[torch.Tensor, torch.Tensor]]: + mapping = json.loads((model / "model.safetensors.index.json").read_text())[ + "weight_map" + ] + + def get(name: str) -> torch.Tensor: + with safe_open(model / mapping[name], framework="pt", device="cpu") as weights: + return weights.get_tensor(name).half() + + result = [] + for layer in range(48): + for role in ("attn", "mlp"): + prefix = f"model.language_model.layers.{layer}.{role}_hyper_connection." + down = torch.zeros((336, 10240), dtype=torch.float16) + down[:320].copy_(get(prefix + "input_mix_weight_down.weight")) + down[320:324].copy_(get(prefix + "block_inject_weight.weight")) + up = get(prefix + "input_mix_weight_up.weight") + result.append((down.cuda(), up.cuda())) + return result + + +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("--quality-inputs", type=int, default=16) + parser.add_argument("--warmup", type=int, default=1000) + parser.add_argument("--replays", type=int, default=150) + parser.add_argument("--stress-replays", type=int, default=32) + args = parser.parse_args() + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + if int(os.environ["WORLD_SIZE"]) != 4 or torch.cuda.get_device_capability() != ( + 7, + 0, + ): + raise RuntimeError("This benchmark requires exactly four SM70 GPUs") + dist.init_process_group("nccl") + group = dist.new_group(backend="gloo") + comm = CustomAllreduce(group=group, device=local_rank, max_size=8 * 1024 * 1024) + try: + if not comm.supports_sm70_qwen38_hc_output_allgather(): + raise RuntimeError("Load the source-matched custom-AR extension") + weights = load_weights(args.model) + generator = torch.Generator(device="cuda").manual_seed(20260905) + xs = torch.randn( + MODULES, 1, 10240, device="cuda", dtype=torch.float16, generator=generator + ) + if not comm.can_sm70_qwen38_hc_shard(xs[0]): + raise RuntimeError("The exact TP4 HC route is unavailable") + sum_generator = torch.Generator(device="cuda").manual_seed(20260905 + rank) + sum_a = torch.randn( + MODULES, 2560, device="cuda", dtype=torch.float16, generator=sum_generator + ) + sum_b = torch.randn_like(sum_a) + peer_sums = [torch.empty_like(sum_a) for _ in range(4)] + dist.all_gather(peer_sums, sum_a + sum_b) + expected_sum = torch.zeros_like(sum_a, dtype=torch.float32) + for peer in peer_sums: + expected_sum.add_(peer.float()) + expected_sum = expected_sum.half() + tp_group = SimpleNamespace(device_communicator=SimpleNamespace(ca_comm=comm)) + + def capture(hidden: bool, overlap: bool = False): + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + outputs, sums = [], [] + aux = torch.cuda.Stream() if overlap else None + with ( + patch( + "vllm.distributed.parallel_state.get_tp_group", + return_value=tp_group, + ), + patch.object( + comm, + "supports_sm70_qwen38_hc_output_allgather", + return_value=hidden, + ), + comm.capture(), + torch.cuda.graph(graph), + ): + main_stream = torch.cuda.current_stream() + if aux is not None: + aux.wait_stream(main_stream) + for i, (down, up) in enumerate(weights): + outputs.extend( + torch.ops.vllm.qwen38_sm70_fp16_fused_hc(xs[i], down, up) + ) + if aux is not None: + with torch.cuda.stream(aux): + sums.append(comm.all_reduce_sum2(sum_a[i], sum_b[i])) + if aux is not None: + main_stream.wait_stream(aux) + torch.cuda.synchronize() + dist.barrier() + return graph, outputs, sums + + # Keep this legacy gate-vs-hidden benchmark on its named routes even + # when the loaded extension also provides the newer fused up path. + fused_override = patch.object( + comm, "supports_sm70_qwen38_hc_up_mix_allgather", return_value=False + ) + with fused_override: + graphs = { + "control": capture(False), + "hidden": capture(True), + "hidden_aux": capture(True, overlap=True), + } + mismatches = {"hidden": 0, "hidden_aux": 0, "sum2_aux": 0} + for case in range(args.quality_inputs): + xs.normal_(generator=generator) + if case == 0: + xs.zero_() + elif case == 1: + xs.mul_(0.01) + for mode, (graph, _, _) in graphs.items(): + # Changing inputs plus repeated epoch wrap tests the HC/MoE + # channels together, not just a single frozen graph replay. + for _ in range(args.stress_replays if mode == "hidden_aux" else 1): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + for mode in ("hidden", "hidden_aux"): + for expected, actual in zip( + graphs["control"][1], graphs[mode][1], strict=True + ): + mismatches[mode] += int( + torch.count_nonzero( + expected.view(torch.int16) != actual.view(torch.int16) + ) + ) + mismatches["sum2_aux"] += int( + torch.count_nonzero( + expected_sum.view(torch.int16) + != torch.stack(graphs["hidden_aux"][2]).view(torch.int16) + ) + ) + quality = [None] * 4 + dist.all_gather_object( + quality, {"rank": rank, "mismatches": mismatches}, group=group + ) + if rank == 0: + print(json.dumps({"quality": quality}), flush=True) + if any(any(q["mismatches"].values()) for q in quality): + if rank == 0: + args.out.write_text(json.dumps({"quality": quality}, indent=2) + "\n") + raise RuntimeError("Production HC or concurrent sum2 is not bitwise") + + # Warm all devices out of idle clocks before paired graph timings. + for mode in ("control", "hidden"): + for _ in range(args.warmup): + graphs[mode][0].replay() + torch.cuda.synchronize() + dist.barrier() + samples = {"control": [], "hidden": []} + for repeat in range(3): + for mode in list(samples) if repeat % 2 == 0 else list(samples)[::-1]: + graph = graphs[mode][0] + for _ in range(20): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(args.replays): + graph.replay() + end.record() + end.synchronize() + times = [None] * 4 + dist.all_gather_object( + times, start.elapsed_time(end) / args.replays, group=group + ) + samples[mode].append(max(times)) + if rank == 0: + medians = {mode: median(values) for mode, values in samples.items()} + result = { + "modules": MODULES, + "includes_combine_norm": False, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "quality": quality, + "quality_inputs": args.quality_inputs, + "aux_stress_replays": args.quality_inputs * args.stress_replays, + "samples_ms": samples, + "median_ms": medians, + "saved_ms": medians["control"] - medians["hidden"], + } + args.out.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2), flush=True) + finally: + comm.close() + dist.destroy_process_group(group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index e55f515ebb..5e246e1180 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -87,6 +87,388 @@ bool _is_weak_contiguous(torch::Tensor& t) { t.numel() * t.element_size()); } +#if !defined(USE_ROCM) +namespace vllm { + +constexpr int kQwen38HcDownLocalElements = 88; +constexpr int kQwen38HcDownLiveElements = 84; +constexpr int kQwen38HcDownLocalLoraElements = 80; +constexpr int kQwen38HcDownLocalInjectionElements = 1; +constexpr int kQwen38HcDownLocalPaddingElements = 3; +constexpr int kQwen38HcDownGatheredElements = + kQwen38HcDownLiveElements * kSm70Tp4PushAllreduceWorldSize; +constexpr int kQwen38HcGateLocalElements = 2560; +constexpr int kQwen38HcGateGatheredElements = + kQwen38HcGateLocalElements * kSm70Tp4PushAllreduceWorldSize; +constexpr int kQwen38HcOutputLocalElements = + kQwen38HcGateLocalElements / kSm70Tp4PushAllreduceWorldSize; + +template +__global__ void __launch_bounds__(128, 1) + sm70_qwen38_hc_down_push_allgather(RankData push_buffers, + const half* __restrict__ input, + half* __restrict__ output) { + static_assert(ngpus == kSm70Tp4PushAllreduceWorldSize); + using P = typename packed_t::P; + constexpr int kElementsPerPack = P::size; + constexpr int kPackedElements = kQwen38HcDownLocalElements / kElementsPerPack; + constexpr int kPackedStride = kSm70Qwen38HcDownPushBytes / sizeof(P); + static_assert(kPackedElements <= kPackedStride); + + auto* local_storage = + const_cast(reinterpret_cast(push_buffers.ptrs[Rank])); + auto* local_epochs = reinterpret_cast( + local_storage + kSm70Qwen38HcPushSignalOffset); + const uint32_t epoch = local_epochs[kSm70Qwen38HcDownEpochIndex]; + const int epoch_offset = epoch * ngpus * kPackedStride; + const int offset = threadIdx.x; + + if (offset < kPackedElements) { + P value = reinterpret_cast(input)[offset]; + #pragma unroll + for (int element = 0; element < P::size; ++element) { + sm70_push_escape_sentinel(value.data[element]); + } + + #pragma unroll + for (int destination_rank = 0; destination_rank < ngpus; + ++destination_rank) { + if (destination_rank == Rank) continue; + auto* destination_base = const_cast( + reinterpret_cast(push_buffers.ptrs[destination_rank])); + void* destination = destination_base + kSm70Qwen38HcDownPushOffset + + (epoch_offset + Rank * kPackedStride) * sizeof(P); + sm70_push_store_volatile_16b(value, destination, offset); + } + + P peer_values[ngpus]; + peer_values[Rank] = value; + while (true) { + bool has_empty_slot = false; + #pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + if (source_rank == Rank) continue; + const void* source = + local_storage + kSm70Qwen38HcDownPushOffset + + (epoch_offset + source_rank * kPackedStride) * sizeof(P); + sm70_push_load_volatile_16b(peer_values[source_rank], source, offset); + #pragma unroll + for (int element = 0; element < P::size; ++element) { + has_empty_slot |= + sm70_push_is_sentinel(peer_values[source_rank].data[element]); + } + } + if (!has_empty_slot) break; + } + + #pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + #pragma unroll + for (int element = 0; element < P::size; ++element) { + const int local_element = offset * kElementsPerPack + element; + if (local_element < kQwen38HcDownLocalLoraElements) { + output[source_rank * kQwen38HcDownLocalLoraElements + local_element] = + peer_values[source_rank].data[element]; + } else if (local_element == kQwen38HcDownLocalLoraElements) { + output[ngpus * kQwen38HcDownLocalLoraElements + source_rank] = + peer_values[source_rank].data[element]; + } else if (local_element < kQwen38HcDownLiveElements) { + const int local_padding = local_element - + kQwen38HcDownLocalLoraElements - + kQwen38HcDownLocalInjectionElements; + output[ngpus * (kQwen38HcDownLocalLoraElements + + kQwen38HcDownLocalInjectionElements) + + source_rank * kQwen38HcDownLocalPaddingElements + + local_padding] = peer_values[source_rank].data[element]; + } + } + } + + P empty; + #pragma unroll + for (int element = 0; element < P::size; ++element) { + *reinterpret_cast(&empty.data[element]) = + kSm70Tp4PushAllreduceSentinel; + } + #pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + if (source_rank == Rank) continue; + void* source = local_storage + kSm70Qwen38HcDownPushOffset + + (epoch_offset + source_rank * kPackedStride) * sizeof(P); + sm70_push_store_volatile_16b(empty, source, offset); + } + } + + __syncthreads(); + if (threadIdx.x == 0) { + local_epochs[kSm70Qwen38HcDownEpochIndex] = + (epoch + 1) % kSm70Tp4PushAllreduceEpochs; + } +} + +DINLINE float qwen38_hc_sigmoid_fp32(float value) { + constexpr uint32_t kLog2E = 0x3fb8aa3b; + const float log2e = __uint_as_float(kLog2E); + const float negated = __fsub_rn(0.0f, value); + const float exponent = __fmul_rn(negated, log2e); + float exp2; + asm volatile("ex2.approx.f32 %0, %1;" : "=f"(exp2) : "f"(exponent)); + const float denominator = __fadd_rn(exp2, 1.0f); + float result; + asm volatile("div.full.f32 %0, %1, %2;" + : "=f"(result) + : "f"(1.0f), "f"(denominator)); + return result; +} + +DINLINE float qwen38_hc_divide_by_count(float value) { + float result; + asm volatile("div.full.f32 %0, %1, %2;" + : "=f"(result) + : "f"(value), "f"(4.0f)); + return result; +} + +// Stream-ordered TP4 HC calls share per-CTA counters. Cooperative launch +// keeps every CTA eligible to make progress while polling its remote peers. +// Pack an exact FP16 output and its 16-bit generation into one aligned word: +// readiness never escapes or changes a floating-point value (including NaNs). +// Two slots are sufficient: a rank cannot produce generation g+2 until every +// peer has produced g+1, hence completed its reads of g. Tag wrap is safe for +// the same reason; only adjacent generations can be in flight. +__device__ __forceinline__ uint4 qwen38_hc_load8(const half* p) { + uint4 v; + asm volatile("ld.global.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p)); + return v; +} +__device__ __forceinline__ float qwen38_hc_half_at(uint4 v, int i) { + const uint32_t word = i < 2 ? v.x : i < 4 ? v.y : i < 6 ? v.z : v.w; + return __half2float(__ushort_as_half((word >> ((i & 1) * 16)) & 0xffffu)); +} + +__global__ void __launch_bounds__(256) + sm70_qwen38_hc_up_mix_push(const half* lora, const half* weight, + const half* branches, half* output, + RankData peers, int rank) { + constexpr int Hidden = 4; + // Constant parameter indices avoid materializing RankData in local memory. + const void* local_peer = rank == 0 ? peers.ptrs[0] + : rank == 1 ? peers.ptrs[1] + : rank == 2 ? peers.ptrs[2] + : peers.ptrs[3]; + auto* local = const_cast(reinterpret_cast(local_peer)); + auto* counters = + reinterpret_cast(local + kSm70Qwen38HcUpFusedEpochOffset); + __shared__ float gates[Hidden * 4]; + __shared__ float partial[Hidden * 4][2]; + const int t = threadIdx.x; + const int lane = t & 31, warp = t >> 5; + const int pair = warp >> 1, kg = warp & 1, kp = kg * 32 + lane; + uint4 lora_values; + if (kp < 40) lora_values = qwen38_hc_load8(lora + kp * 8); + // Same 8-term FMA chain, XOR tree, cross-warp add and FP16 gate boundary + // as the accepted Triton up projection. Only the row assignment changes. + #pragma unroll + for (int group = 0; group < Hidden / 2; ++group) { + const int a = group * 8 + pair, b = a + 4; + const int ra = (a % 4) * 2560 + rank * 640 + blockIdx.x * Hidden + a / 4; + const int rb = (b % 4) * 2560 + rank * 640 + blockIdx.x * Hidden + b / 4; + float va = 0.f, vb = 0.f; + if (kp < 40) { + const int k = kp * 8; + const uint4 wa = qwen38_hc_load8(weight + ra * 320 + k); + const uint4 wb = qwen38_hc_load8(weight + rb * 320 + k); + const float x1 = qwen38_hc_half_at(lora_values, 1); + va = __fmul_rn(x1, qwen38_hc_half_at(wa, 1)); + vb = __fmul_rn(x1, qwen38_hc_half_at(wb, 1)); + va = __fmaf_rn(qwen38_hc_half_at(lora_values, 0), + qwen38_hc_half_at(wa, 0), va); + vb = __fmaf_rn(qwen38_hc_half_at(lora_values, 0), + qwen38_hc_half_at(wb, 0), vb); + #pragma unroll + for (int e = 2; e < 8; ++e) { + const float x = qwen38_hc_half_at(lora_values, e); + va = __fmaf_rn(x, qwen38_hc_half_at(wa, e), va); + vb = __fmaf_rn(x, qwen38_hc_half_at(wb, e), vb); + } + } + #pragma unroll + for (int d = 16; d > 0; d >>= 1) { + va = __fadd_rn(va, __shfl_xor_sync(0xffffffff, va, d)); + vb = __fadd_rn(vb, __shfl_xor_sync(0xffffffff, vb, d)); + } + if (lane == 0) { + partial[a][kg] = va; + partial[b][kg] = vb; + } + } + __syncthreads(); + if (t < Hidden * 4) + gates[t] = qwen38_hc_sigmoid_fp32( + __half2float(__float2half_rn(__fadd_rn(partial[t][0], partial[t][1])))); + __syncthreads(); + if (t < Hidden) { + const int h = blockIdx.x * Hidden + t; + float mixed = 0.f; + #pragma unroll + for (int branch = 0; branch < 4; ++branch) + mixed = __fmaf_rn(gates[t * 4 + branch], + __half2float(branches[branch * 2560 + rank * 640 + h]), + mixed); + float scaled; + asm("div.full.f32 %0, %1, %2;" : "=f"(scaled) : "f"(mixed), "f"(4.f)); + const half value = __float2half_rn(scaled); + { + const uint32_t generation = counters[blockIdx.x] + 1u; + const uint32_t tag = generation & 0xffffu; + const uint32_t packet = (tag << 16) | __half_as_ushort(value); + const int slot = (generation & 1u) * 4 * 640; + #pragma unroll + for (int dest = 0; dest < 4; ++dest) { + if (dest == rank) continue; + auto* p = reinterpret_cast( + const_cast( + reinterpret_cast(peers.ptrs[dest])) + + kSm70Qwen38HcUpFusedPacketOffset) + + slot + rank * 640 + h; + asm volatile("st.volatile.global.u32 [%0], %1;" ::"l"(p), "r"(packet) + : "memory"); + } + output[rank * 640 + h] = value; + #pragma unroll + for (int src = 0; src < 4; ++src) { + if (src == rank) continue; + auto* p = reinterpret_cast( + local + kSm70Qwen38HcUpFusedPacketOffset) + + slot + src * 640 + h; + uint32_t received; + do { + asm volatile("ld.volatile.global.u32 %0, [%1];" + : "=r"(received) + : "l"(p) + : "memory"); + } while ((received >> 16) != tag); + output[src * 640 + h] = __ushort_as_half(received & 0xffffu); + } + } + } + { + __syncthreads(); + if (t == 0) counters[blockIdx.x] += 1; + } +} + +template +__global__ void __launch_bounds__(512, 1) + sm70_qwen38_hc_gate_push_mix(RankData push_buffers, + const half* __restrict__ local_gate, + const half* __restrict__ branches, + half* __restrict__ output, + int packed_elements) { + static_assert(ngpus == kSm70Tp4PushAllreduceWorldSize); + using P = typename packed_t::P; + constexpr int kPackedStride = kSm70Qwen38HcGatePushBytes / sizeof(P); + + auto* local_storage = + const_cast(reinterpret_cast(push_buffers.ptrs[Rank])); + auto* local_epochs = reinterpret_cast( + local_storage + kSm70Qwen38HcPushSignalOffset); + const int epoch_index = kSm70Qwen38HcGateEpochIndexBase + blockIdx.x; + const uint32_t epoch = local_epochs[epoch_index]; + const int epoch_offset = epoch * ngpus * kPackedStride; + const int offset = blockIdx.x * blockDim.x + threadIdx.x; + + if (offset < packed_elements) { + P value = reinterpret_cast(local_gate)[offset]; + #pragma unroll + for (int element = 0; element < P::size; ++element) { + sm70_push_escape_sentinel(value.data[element]); + } + + #pragma unroll + for (int destination_rank = 0; destination_rank < ngpus; + ++destination_rank) { + if (destination_rank == Rank) continue; + auto* destination_base = const_cast( + reinterpret_cast(push_buffers.ptrs[destination_rank])); + void* destination = destination_base + kSm70Qwen38HcGatePushOffset + + (epoch_offset + Rank * kPackedStride) * sizeof(P); + sm70_push_store_volatile_16b(value, destination, offset); + } + + P peer_values[ngpus]; + peer_values[Rank] = value; + while (true) { + bool has_empty_slot = false; + #pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + if (source_rank == Rank) continue; + const void* source = + local_storage + kSm70Qwen38HcGatePushOffset + + (epoch_offset + source_rank * kPackedStride) * sizeof(P); + sm70_push_load_volatile_16b(peer_values[source_rank], source, offset); + #pragma unroll + for (int element = 0; element < P::size; ++element) { + has_empty_slot |= + sm70_push_is_sentinel(peer_values[source_rank].data[element]); + } + } + if (!has_empty_slot) break; + } + + if constexpr (GatherOutput) { + // Up already mixed all branches for 640 hidden coordinates. Gather + // these final FP16 values, without a second arithmetic/rounding step. + // Reuse the isolated HC gate channel and its existing epoch protocol; + // the MoE/shared-expert stream uses a separate channel. + #pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + reinterpret_cast( + output + source_rank * kQwen38HcOutputLocalElements)[offset] = + peer_values[source_rank]; + } + } else { + #pragma unroll + for (int element = 0; element < P::size; ++element) { + const int hidden = offset * P::size + element; + float result = 0.0f; + #pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + const float gate = __half2float(peer_values[source_rank].data[element]); + const float branch = __half2float( + branches[source_rank * kQwen38HcGateLocalElements + hidden]); + result = __fmaf_rn(qwen38_hc_sigmoid_fp32(gate), branch, result); + } + output[hidden] = __float2half_rn(qwen38_hc_divide_by_count(result)); + } + } + + P empty; + #pragma unroll + for (int element = 0; element < P::size; ++element) { + *reinterpret_cast(&empty.data[element]) = + kSm70Tp4PushAllreduceSentinel; + } + #pragma unroll + for (int source_rank = 0; source_rank < ngpus; ++source_rank) { + if (source_rank == Rank) continue; + void* source = local_storage + kSm70Qwen38HcGatePushOffset + + (epoch_offset + source_rank * kPackedStride) * sizeof(P); + sm70_push_store_volatile_16b(empty, source, offset); + } + } + + __syncthreads(); + if (threadIdx.x == 0) { + local_epochs[epoch_index] = (epoch + 1) % kSm70Tp4PushAllreduceEpochs; + } +} + +} // namespace vllm +#endif + /** * Performs an out-of-place allreduce and stores result in out. * @@ -413,6 +795,172 @@ void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, } } +void sm70_qwen38_hc_down_allgather(fptr_t _fa, torch::Tensor& input, + torch::Tensor& output) { +#if defined(USE_ROCM) + TORCH_CHECK(false, "SM70 Qwen3.8 HC all-gather is unavailable on ROCm"); +#else + auto fa = reinterpret_cast(_fa); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + TORCH_CHECK_EQ(fa->world_size_, vllm::kSm70Tp4PushAllreduceWorldSize); + TORCH_CHECK(fa->fully_connected_ && fa->sm70_tp4_push_buffers_registered_); + TORCH_CHECK_EQ(input.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(output.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(input.numel(), vllm::kQwen38HcDownLocalElements); + TORCH_CHECK_EQ(output.numel(), vllm::kQwen38HcDownGatheredElements); + TORCH_CHECK(_is_weak_contiguous(input) && _is_weak_contiguous(output)); + #define VLLM_LAUNCH_QWEN38_HC_DOWN(RANK) \ + vllm::sm70_qwen38_hc_down_push_allgather<4, RANK><<<1, 32, 0, stream>>>( \ + fa->sm70_tp4_push_buffers_, \ + reinterpret_cast(input.data_ptr()), \ + reinterpret_cast(output.data_ptr())) + switch (fa->rank_) { + case 0: + VLLM_LAUNCH_QWEN38_HC_DOWN(0); + break; + case 1: + VLLM_LAUNCH_QWEN38_HC_DOWN(1); + break; + case 2: + VLLM_LAUNCH_QWEN38_HC_DOWN(2); + break; + default: + VLLM_LAUNCH_QWEN38_HC_DOWN(3); + break; + } + #undef VLLM_LAUNCH_QWEN38_HC_DOWN +#endif +} + +void sm70_qwen38_hc_gate_mix(fptr_t _fa, torch::Tensor& local_gate, + torch::Tensor& branches, torch::Tensor& output) { +#if defined(USE_ROCM) + TORCH_CHECK(false, "SM70 Qwen3.8 HC gate-mix is unavailable on ROCm"); +#else + auto fa = reinterpret_cast(_fa); + const at::cuda::OptionalCUDAGuard device_guard(device_of(local_gate)); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + TORCH_CHECK_EQ(fa->world_size_, vllm::kSm70Tp4PushAllreduceWorldSize); + TORCH_CHECK(fa->fully_connected_ && fa->sm70_tp4_push_buffers_registered_); + TORCH_CHECK_EQ(local_gate.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(branches.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(output.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(local_gate.numel(), vllm::kQwen38HcGateLocalElements); + TORCH_CHECK_EQ(branches.numel(), vllm::kQwen38HcGateGatheredElements); + TORCH_CHECK_EQ(output.numel(), vllm::kQwen38HcGateLocalElements); + TORCH_CHECK(_is_weak_contiguous(local_gate) && + _is_weak_contiguous(branches) && _is_weak_contiguous(output)); + constexpr int kPackedElements = + vllm::kQwen38HcGateLocalElements / vllm::packed_t::P::size; + constexpr int kThreads = 32; + constexpr int kBlocks = (kPackedElements + kThreads - 1) / kThreads; + #define VLLM_LAUNCH_QWEN38_HC_GATE(RANK) \ + vllm::sm70_qwen38_hc_gate_push_mix<4, RANK> \ + <<>>( \ + fa->sm70_tp4_push_buffers_, \ + reinterpret_cast(local_gate.data_ptr()), \ + reinterpret_cast(branches.data_ptr()), \ + reinterpret_cast(output.data_ptr()), kPackedElements) + switch (fa->rank_) { + case 0: + VLLM_LAUNCH_QWEN38_HC_GATE(0); + break; + case 1: + VLLM_LAUNCH_QWEN38_HC_GATE(1); + break; + case 2: + VLLM_LAUNCH_QWEN38_HC_GATE(2); + break; + default: + VLLM_LAUNCH_QWEN38_HC_GATE(3); + break; + } + #undef VLLM_LAUNCH_QWEN38_HC_GATE +#endif +} + +void sm70_qwen38_hc_up_mix_allgather(fptr_t _fa, torch::Tensor& lora, + torch::Tensor& weight, + torch::Tensor& branches, + torch::Tensor& output) { +#if defined(USE_ROCM) + TORCH_CHECK(false, "SM70 Qwen3.8 HC up/mix is unavailable on ROCm"); +#else + TORCH_CHECK(lora.is_cuda()); + const at::cuda::OptionalCUDAGuard device_guard(device_of(lora)); + for (const auto* tensor : {&lora, &weight, &branches, &output}) { + TORCH_CHECK(tensor->device() == lora.device()); + TORCH_CHECK(tensor->scalar_type() == at::ScalarType::Half); + TORCH_CHECK(tensor->is_contiguous()); + } + TORCH_CHECK(lora.numel() == 336 && weight.numel() == 10240 * 320 && + branches.numel() == 10240 && output.numel() == 2560); + TORCH_CHECK(reinterpret_cast(lora.data_ptr()) % 16 == 0 && + reinterpret_cast(weight.data_ptr()) % 16 == 0); + auto* fa = reinterpret_cast(_fa); + TORCH_CHECK(fa->world_size_ == 4 && fa->fully_connected_ && + fa->sm70_tp4_push_buffers_registered_); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + const half* lp = reinterpret_cast(lora.data_ptr()); + const half* wp = reinterpret_cast(weight.data_ptr()); + const half* xp = reinterpret_cast(branches.data_ptr()); + half* out = reinterpret_cast(output.data_ptr()); + auto peers = fa->sm70_tp4_push_buffers_; + int rank = fa->rank_; + void* args[] = {&lp, &wp, &xp, &out, &peers, &rank}; + CUDACHECK(cudaLaunchCooperativeKernel( + reinterpret_cast(vllm::sm70_qwen38_hc_up_mix_push), + dim3(vllm::kSm70Qwen38HcUpFusedBlocks), dim3(256), args, 0, stream)); +#endif +} + +void sm70_qwen38_hc_output_allgather(fptr_t _fa, torch::Tensor& local_block, + torch::Tensor& output) { +#if defined(USE_ROCM) + TORCH_CHECK(false, "SM70 Qwen3.8 HC output all-gather is unavailable on ROCm"); +#else + auto fa = reinterpret_cast(_fa); + TORCH_CHECK(local_block.is_cuda() && output.is_cuda()); + TORCH_CHECK(local_block.device() == output.device()); + const at::cuda::OptionalCUDAGuard device_guard(device_of(local_block)); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + TORCH_CHECK_EQ(fa->world_size_, vllm::kSm70Tp4PushAllreduceWorldSize); + TORCH_CHECK(fa->fully_connected_ && fa->sm70_tp4_push_buffers_registered_); + TORCH_CHECK_EQ(local_block.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(output.scalar_type(), at::ScalarType::Half); + TORCH_CHECK_EQ(local_block.numel(), vllm::kQwen38HcOutputLocalElements); + TORCH_CHECK_EQ(output.numel(), vllm::kQwen38HcGateLocalElements); + TORCH_CHECK(local_block.is_contiguous() && output.is_contiguous()); + constexpr int kPackedElements = + vllm::kQwen38HcOutputLocalElements / vllm::packed_t::P::size; + constexpr int kThreads = 32; + constexpr int kBlocks = (kPackedElements + kThreads - 1) / kThreads; + #define VLLM_LAUNCH_QWEN38_HC_OUTPUT(RANK) \ + vllm::sm70_qwen38_hc_gate_push_mix<4, RANK, true> \ + <<>>( \ + fa->sm70_tp4_push_buffers_, \ + reinterpret_cast(local_block.data_ptr()), \ + nullptr, reinterpret_cast(output.data_ptr()), \ + kPackedElements) + switch (fa->rank_) { + case 0: + VLLM_LAUNCH_QWEN38_HC_OUTPUT(0); + break; + case 1: + VLLM_LAUNCH_QWEN38_HC_OUTPUT(1); + break; + case 2: + VLLM_LAUNCH_QWEN38_HC_OUTPUT(2); + break; + default: + VLLM_LAUNCH_QWEN38_HC_OUTPUT(3); + break; + } + #undef VLLM_LAUNCH_QWEN38_HC_OUTPUT +#endif +} + void top1_argmax(fptr_t _fa, torch::Tensor& input_pair, torch::Tensor& output, fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) { auto fa = reinterpret_cast(_fa); diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 66e6765ff3..99d7dce2c2 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -75,10 +75,43 @@ constexpr size_t kSm70Tp4PushAllreduceQwen4ExpMtp5Bytes = 5 * 2560 * sizeof(half); constexpr size_t kSm70Tp4PushAllreduceSignalBytes = ((kSm70Tp4PushAllreduceBlocks * sizeof(uint32_t) + 127) / 128) * 128; -constexpr size_t kSm70Tp4PushAllreduceBufferBytes = +constexpr size_t kSm70Tp4PushAllreduceGenericBufferBytes = kSm70Tp4PushAllreduceSignalBytes + kSm70Tp4PushAllreduceEpochs * kSm70Tp4PushAllreduceWorldSize * kSm70Tp4PushAllreduceBytes; +// HC decode can overlap the ordinary MoE push collective on vLLM's auxiliary +// stream. Keep both its epoch words and payloads disjoint so an HC poll cannot +// observe or clear a concurrently running all-reduce packet. The ordinary +// collective layout above remains unchanged. +constexpr int kSm70Qwen38HcGatePushBlocks = 10; +constexpr int kSm70Qwen38HcDownEpochIndex = 0; +constexpr int kSm70Qwen38HcGateEpochIndexBase = 1; +constexpr size_t kSm70Qwen38HcPushSignalOffset = + kSm70Tp4PushAllreduceGenericBufferBytes; +constexpr size_t kSm70Qwen38HcPushSignalBytes = 128; +constexpr size_t kSm70Qwen38HcDownPushBytes = 256; +constexpr size_t kSm70Qwen38HcGatePushBytes = 2560 * sizeof(half); +constexpr size_t kSm70Qwen38HcDownPushOffset = + kSm70Qwen38HcPushSignalOffset + kSm70Qwen38HcPushSignalBytes; +constexpr size_t kSm70Qwen38HcGatePushOffset = + kSm70Qwen38HcDownPushOffset + kSm70Tp4PushAllreduceEpochs * + kSm70Tp4PushAllreduceWorldSize * + kSm70Qwen38HcDownPushBytes; +constexpr size_t kSm70Qwen38HcUpFusedEpochOffset = + kSm70Qwen38HcGatePushOffset + kSm70Tp4PushAllreduceEpochs * + kSm70Tp4PushAllreduceWorldSize * + kSm70Qwen38HcGatePushBytes; +// The fused up/mix/gather uses 160 independent generation counters and exact +// half-plus-tag packets, separate from both legacy HC and auxiliary MoE data. +constexpr int kSm70Qwen38HcUpFusedBlocks = 160; +constexpr size_t kSm70Qwen38HcUpFusedPacketOffset = + kSm70Qwen38HcUpFusedEpochOffset + + kSm70Qwen38HcUpFusedBlocks * sizeof(uint32_t); +constexpr size_t kSm70Tp4PushAllreduceBufferBytes = + kSm70Qwen38HcUpFusedPacketOffset + + kSm70Tp4PushAllreduceEpochs * 4 * 640 * sizeof(uint32_t); +static_assert(kSm70Qwen38HcGateEpochIndexBase + kSm70Qwen38HcGatePushBlocks <= + kSm70Qwen38HcPushSignalBytes / sizeof(uint32_t)); inline int sm70_tp4_push_allreduce_blocks(size_t bytes) { if (bytes == kSm70Tp4PushAllreduceBytes) { @@ -1505,11 +1538,25 @@ class CustomAllreduce { } sm70_tp4_push_buffers_.ptrs[peer] = ptrs[peer]; } - auto* local_data = + auto* generic_data = static_cast(ptrs[rank_]) + kSm70Tp4PushAllreduceSignalBytes; + CUDACHECK(cudaMemset(generic_data, kSm70Tp4PushAllreduceSentinelByte, + kSm70Tp4PushAllreduceGenericBufferBytes - + kSm70Tp4PushAllreduceSignalBytes)); + auto* hc_signal = + static_cast(ptrs[rank_]) + kSm70Qwen38HcPushSignalOffset; + CUDACHECK(cudaMemset(hc_signal, 0, kSm70Qwen38HcPushSignalBytes)); + auto* hc_data = + static_cast(ptrs[rank_]) + kSm70Qwen38HcDownPushOffset; + CUDACHECK(cudaMemset( + hc_data, kSm70Tp4PushAllreduceSentinelByte, + kSm70Qwen38HcUpFusedEpochOffset - kSm70Qwen38HcDownPushOffset)); + // The first fused packet uses generation 1; zero is initially invalid. + auto* hc_up = + static_cast(ptrs[rank_]) + kSm70Qwen38HcUpFusedEpochOffset; CUDACHECK(cudaMemset( - local_data, kSm70Tp4PushAllreduceSentinelByte, - kSm70Tp4PushAllreduceBufferBytes - kSm70Tp4PushAllreduceSignalBytes)); + hc_up, 0, + kSm70Tp4PushAllreduceBufferBytes - kSm70Qwen38HcUpFusedEpochOffset)); sm70_tp4_push_buffers_registered_ = true; } @@ -1862,10 +1909,16 @@ class CustomAllreduce { size /= d; auto bytes = size * sizeof(typename packed_t::P); if constexpr (std::is_same_v) { + const char* qwen4_exp_m1 = + std::getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1"); + const bool qwen4_exp_m1_enabled = + bytes == kSm70Tp4PushAllreduceQwen4ExpBytes && + (qwen4_exp_m1 == nullptr || std::strcmp(qwen4_exp_m1, "1") == 0); if (sm70_tp4_push_buffers_registered_ && status == cudaStreamCaptureStatusActive && world_size_ == kSm70Tp4PushAllreduceWorldSize && fully_connected_ && - bytes == kSm70Tp4PushAllreduceQwen4ExpMtp5Bytes && + (bytes == kSm70Tp4PushAllreduceQwen4ExpMtp5Bytes || + qwen4_exp_m1_enabled) && custom_allreduce_current_device_is_sm70()) { const int push_blocks = sm70_tp4_push_allreduce_blocks(bytes); if (push_blocks > 0) { diff --git a/csrc/ops.h b/csrc/ops.h index 57d4bfb1c6..37cad87e7d 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -391,6 +391,9 @@ void sm70_dynamic_draft_vocab_refresh_tail_weight_out( void sm70_f16_gate_mul_out(torch::Tensor out, torch::Tensor _in_feats, torch::Tensor _gate_weight); +void qwen38_shared_gate_exact_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weight); + int64_t sm70_gemm_import_cache(torch::Tensor device_hint, const std::string& path); @@ -523,6 +526,17 @@ void nvfp4_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor expert_ids, bool broadcast_input, int64_t split_k); +void nvfp4_qwen38_w2_direct_reduce_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, + torch::Tensor scales, + torch::Tensor expert_ids, + torch::Tensor topk_weights); + +void nvfp4_qwen38_w13_fused_swiglu_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, + torch::Tensor scales, + torch::Tensor expert_ids); + void nvfp4_moe_qpn_mtp5_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor weights, torch::Tensor scales, torch::Tensor expert_ids, bool broadcast_input, @@ -644,6 +658,17 @@ 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_down_allgather(fptr_t _fa, torch::Tensor& input, + torch::Tensor& output); +void sm70_qwen38_hc_gate_mix(fptr_t _fa, torch::Tensor& local_gate, + torch::Tensor& branches, torch::Tensor& output); +void sm70_qwen38_hc_output_allgather(fptr_t _fa, torch::Tensor& local_block, + torch::Tensor& output); + +void sm70_qwen38_hc_up_mix_allgather(fptr_t _fa, torch::Tensor& lora, + torch::Tensor& weight, + torch::Tensor& branches, + torch::Tensor& output); void top1_argmax(fptr_t _fa, torch::Tensor& input_pair, torch::Tensor& output, fptr_t reg_buffer, int64_t reg_buffer_sz_bytes); void tile_runtime_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out, diff --git a/csrc/qsa_lexicographic_topk.cuh b/csrc/qsa_lexicographic_topk.cuh index 9ced91d081..3fe11052ca 100644 --- a/csrc/qsa_lexicographic_topk.cuh +++ b/csrc/qsa_lexicographic_topk.cuh @@ -13,6 +13,7 @@ namespace vllm::qsa { constexpr int kLexicographicTopKThreads = 1024; constexpr int kLexicographicTopKBins = 256; +constexpr int kLexicographicTopKDecodeCandidateCapacity = 2304; __device__ __forceinline__ uint32_t ordered_float_bits(float value) { // IEEE -0.0 and +0.0 compare equal, so keep them in the same score bucket @@ -37,6 +38,72 @@ struct LexicographicTopKShared { uint32_t chunk_equal_base; }; +template +struct LexicographicDecodeTopKShared { + using BlockScan = cub::BlockScan; + + // The decode fast path first selects one coarse radix bucket, then scans + // only that bucket for the remaining bytes. Keep two buffers so compaction + // never overwrites input indices that another warp has not consumed yet. + uint32_t histogram[2][kLexicographicTopKBins + 128]; + int32_t candidates[2][kLexicographicTopKDecodeCandidateCapacity]; + typename BlockScan::TempStorage scan; + uint32_t prefix; + uint32_t pivot; + uint32_t remaining; + uint32_t remaining_ties; + uint32_t candidate_count[2]; + uint32_t threshold_bin; + uint32_t greater_seen; + uint32_t equal_seen; + uint32_t chunk_greater_base; + uint32_t chunk_equal_base; +}; + +template +__device__ __forceinline__ void decode_suffix_scan_histogram( + LexicographicDecodeTopKShared& shared) { +#pragma unroll + for (int pass = 0; pass < 8; ++pass) { + const int distance = 1 << pass; + const int source = pass & 1; + if (threadIdx.x < kLexicographicTopKBins) { + uint32_t value = shared.histogram[source][threadIdx.x]; + if (threadIdx.x + distance < kLexicographicTopKBins) { + value += shared.histogram[source][threadIdx.x + distance]; + } + shared.histogram[source ^ 1][threadIdx.x] = value; + } + __syncthreads(); + } +} + +template +__device__ __forceinline__ void decode_choose_threshold( + LexicographicDecodeTopKShared& shared, int shift) { + if (threadIdx.x < kLexicographicTopKBins && + shared.histogram[0][threadIdx.x] > shared.remaining && + shared.histogram[0][threadIdx.x + 1] <= shared.remaining) { + shared.threshold_bin = threadIdx.x; + } + __syncthreads(); + if (threadIdx.x == 0) { + const uint32_t bin = shared.threshold_bin; + const uint32_t greater = shared.histogram[0][bin + 1]; + shared.remaining -= greater; + shared.prefix |= bin << shift; + if (shared.remaining == 0) { + const uint32_t low_mask = shift == 0 ? 0u : ((uint32_t{1} << shift) - 1); + shared.pivot = shared.prefix | low_mask; + shared.remaining_ties = 0; + } else if (shift == 0) { + shared.pivot = shared.prefix; + shared.remaining_ties = shared.remaining; + } + } + __syncthreads(); +} + template __global__ __launch_bounds__(kLexicographicTopKThreads) void qsa_lexicographic_topk_kernel( @@ -153,14 +220,196 @@ __launch_bounds__(kLexicographicTopKThreads) void qsa_lexicographic_topk_kernel( } } +// Single-token QSA decode has only about two thousand live block scores at the +// common 8K context length. After the first radix byte, scanning all scores for +// the other three bytes wastes most of the work. Compact the selected coarse +// bucket into shared memory and refine that much smaller set instead. Integer +// counters and the final increasing-index pass retain exact tie-breaking. +template +__global__ +__launch_bounds__(kLexicographicTopKThreads) void qsa_lexicographic_decode_topk_kernel( + const float* __restrict__ logits, const int32_t* __restrict__ lengths, + int32_t* __restrict__ output, uint32_t columns) { + const uint32_t tx = threadIdx.x; + const int32_t raw_length = lengths[0]; + const uint32_t length = + raw_length > 0 ? min(static_cast(raw_length), columns) : 0; + + if (length <= TopK) { + for (uint32_t index = tx; index < TopK; + index += kLexicographicTopKThreads) { + output[index] = index < length ? static_cast(index) : -1; + } + return; + } + + __shared__ LexicographicDecodeTopKShared shared; + if (tx == 0) { + shared.prefix = 0; + shared.remaining = TopK; + shared.remaining_ties = 0; + shared.candidate_count[0] = 0; + } + __syncthreads(); + + if (length > kLexicographicTopKDecodeCandidateCapacity) { + // Preserve the original exact four-pass algorithm for long contexts, + // without a host synchronization or a second kernel launch. +#pragma unroll + for (int pass = 0; pass < 4; ++pass) { + for (uint32_t bin = tx; bin < kLexicographicTopKBins; + bin += kLexicographicTopKThreads) { + shared.histogram[0][bin] = 0; + } + __syncthreads(); + + const int shift = 24 - pass * 8; + const uint32_t prefix = shared.prefix; + const uint32_t prefix_mask = + pass == 0 ? 0 : (~uint32_t{0} << (shift + 8)); + for (uint32_t index = tx; index < length; + index += kLexicographicTopKThreads) { + const uint32_t key = ordered_float_bits(logits[index]); + if ((key & prefix_mask) == prefix) { + atomicAdd(&shared.histogram[0][(key >> shift) & 0xffu], 1u); + } + } + __syncthreads(); + + if (tx == 0) { + uint32_t remaining = shared.remaining; + for (int bin = kLexicographicTopKBins - 1; bin >= 0; --bin) { + const uint32_t count = shared.histogram[0][bin]; + if (remaining > count) { + remaining -= count; + } else { + shared.prefix |= static_cast(bin) << shift; + shared.remaining = remaining; + break; + } + } + } + __syncthreads(); + } + if (tx == 0) { + shared.pivot = shared.prefix; + shared.remaining_ties = shared.remaining; + } + __syncthreads(); + } else { + // Coarse pass over the complete score row. + if (tx < kLexicographicTopKBins + 1) shared.histogram[0][tx] = 0; + __syncthreads(); + for (uint32_t index = tx; index < length; + index += kLexicographicTopKThreads) { + const uint32_t key = ordered_float_bits(logits[index]); + atomicAdd(&shared.histogram[0][key >> 24], 1u); + } + __syncthreads(); + decode_suffix_scan_histogram(shared); + decode_choose_threshold(shared, 24); + + if (shared.remaining != 0) { + if (tx < kLexicographicTopKBins + 1) shared.histogram[0][tx] = 0; + __syncthreads(); + for (uint32_t index = tx; index < length; + index += kLexicographicTopKThreads) { + const uint32_t key = ordered_float_bits(logits[index]); + if ((key & 0xff000000u) == shared.prefix) { + const uint32_t position = atomicAdd(&shared.candidate_count[0], 1u); + shared.candidates[0][position] = static_cast(index); + atomicAdd(&shared.histogram[0][(key >> 16) & 0xffu], 1u); + } + } + __syncthreads(); + } + +#pragma unroll + for (int radix_pass = 0; radix_pass < 3; ++radix_pass) { + if (shared.remaining == 0) break; + const int shift = 16 - radix_pass * 8; + decode_suffix_scan_histogram(shared); + decode_choose_threshold(shared, shift); + if (shared.remaining == 0 || shift == 0) break; + + const int source = radix_pass & 1; + const int target = source ^ 1; + if (tx == 0) shared.candidate_count[target] = 0; + if (tx < kLexicographicTopKBins + 1) shared.histogram[0][tx] = 0; + __syncthreads(); + const uint32_t count = shared.candidate_count[source]; + const uint32_t prefix_mask = ~uint32_t{0} << shift; + const int next_shift = shift - 8; + for (uint32_t item = tx; item < count; + item += kLexicographicTopKThreads) { + const int32_t index = shared.candidates[source][item]; + const uint32_t key = ordered_float_bits(logits[index]); + if ((key & prefix_mask) == shared.prefix) { + const uint32_t position = + atomicAdd(&shared.candidate_count[target], 1u); + shared.candidates[target][position] = index; + atomicAdd(&shared.histogram[0][(key >> next_shift) & 0xffu], 1u); + } + } + __syncthreads(); + } + } + + if (tx == 0) { + shared.greater_seen = 0; + shared.equal_seen = 0; + } + __syncthreads(); + + // Emit in original index order, matching QSA's canonical accumulation order. + using BlockScan = typename LexicographicDecodeTopKShared::BlockScan; + for (uint32_t base = 0; base < length; base += kLexicographicTopKThreads) { + const uint32_t index = base + tx; + const uint32_t key = index < length ? ordered_float_bits(logits[index]) : 0; + const uint32_t greater = index < length && key > shared.pivot ? 1u : 0u; + const uint32_t equal = index < length && key == shared.pivot ? 1u : 0u; + const uint64_t counts = (static_cast(greater) << 32) | equal; + uint64_t prefix_counts = 0; + uint64_t aggregate_counts = 0; + BlockScan(shared.scan) + .ExclusiveSum(counts, prefix_counts, aggregate_counts); + __syncthreads(); + if (tx == 0) { + shared.chunk_greater_base = shared.greater_seen; + shared.chunk_equal_base = shared.equal_seen; + shared.greater_seen += static_cast(aggregate_counts >> 32); + shared.equal_seen += static_cast(aggregate_counts); + } + __syncthreads(); + const uint32_t greater_before = + shared.chunk_greater_base + static_cast(prefix_counts >> 32); + const uint32_t equal_before = + shared.chunk_equal_base + static_cast(prefix_counts); + const bool selected = + greater || (equal && equal_before < shared.remaining_ties); + if (selected) { + const uint32_t offset = + greater_before + min(equal_before, shared.remaining_ties); + output[offset] = static_cast(index); + } + __syncthreads(); + } +} + template void launch_qsa_lexicographic_topk(const float* logits, const int32_t* lengths, int32_t* output, uint32_t num_rows, uint32_t columns, uint32_t stride, cudaStream_t stream) { - qsa_lexicographic_topk_kernel - <<>>( - logits, lengths, output, num_rows, columns, stride); + if (num_rows == 1) { + qsa_lexicographic_decode_topk_kernel + <<<1, kLexicographicTopKThreads, 0, stream>>>(logits, lengths, output, + columns); + } else { + qsa_lexicographic_topk_kernel + <<>>( + logits, lengths, output, num_rows, columns, stride); + } } } // namespace vllm::qsa diff --git a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu index ad00901b59..f864dcc7e8 100644 --- a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -11,6 +11,58 @@ namespace { +constexpr int kQwen38SharedGateHidden = 2560; +constexpr int kQwen38SharedGateThreads = 256; + +__device__ __forceinline__ float qwen38_shared_gate_warp_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = __fadd_rn(value, __shfl_down_sync(0xffffffffU, value, offset)); + } + return value; +} + +__global__ void qwen38_shared_gate_exact_kernel( + half* __restrict__ output, const half* __restrict__ input, + const half* __restrict__ weight) { + constexpr int kValuesPerThread = + kQwen38SharedGateHidden / kQwen38SharedGateThreads; + const int tid = threadIdx.x; + float value = 0.0f; +#pragma unroll + for (int item = 0; item < kValuesPerThread; ++item) { + const int index = tid + item * kQwen38SharedGateThreads; + value = __fmaf_rn(__half2float(input[index]), __half2float(weight[index]), + value); + } + value = qwen38_shared_gate_warp_sum(value); + + __shared__ float warp_partials[kQwen38SharedGateThreads / 32]; + __shared__ half shared_gate; + if ((tid & 31) == 0) { + warp_partials[tid >> 5] = value; + } + __syncthreads(); + if (tid < 32) { + value = tid < kQwen38SharedGateThreads / 32 ? warp_partials[tid] : 0.0f; + value = qwen38_shared_gate_warp_sum(value); + if (tid == 0) { + // Preserve the eager FP16 linear and sigmoid materialization points. + const half linear = __float2half_rn(value); + const float rounded_linear = __half2float(linear); + shared_gate = __float2half_rn(1.0f / (1.0f + __expf(-rounded_linear))); + } + } + __syncthreads(); + + const half2 gate = __half2half2(shared_gate); + auto* output2 = reinterpret_cast(output); + for (int index = tid; index < kQwen38SharedGateHidden / 2; + index += blockDim.x) { + output2[index] = __hmul2(output2[index], gate); + } +} + __device__ __forceinline__ void dequant_e2m1x8(unsigned packed, half2 scale, half2 out[4]) { constexpr unsigned kSign = 0x80008000u; @@ -138,7 +190,7 @@ __global__ void mxfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, } } -template +template __global__ void nvfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, const uint32_t* __restrict__ weights, const half* __restrict__ scales, @@ -153,7 +205,12 @@ __global__ void nvfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, const int route = blockIdx.y; const int expert = __ldg(expert_ids + route); if (expert < 0 || expert >= 512) { - if (threadIdx.x < 32) { + if constexpr (kFusedSwiGLU) { + if (threadIdx.x < 16) { + output[static_cast(route) * (n / 2) + tile * 16 + threadIdx.x] = + __float2half(0.0f); + } + } else if (threadIdx.x < 32) { output[static_cast(route) * n + tile * 32 + threadIdx.x] = __float2half(0.0f); } @@ -238,8 +295,120 @@ __global__ void nvfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, for (int k_warp = 0; k_warp < kSplitK; ++k_warp) { value += partials[k_warp][lane]; } - output[static_cast(route) * n + tile * 32 + lane] = - __float2half(value); + const half rounded = __float2half(value); + if constexpr (kFusedSwiGLU) { + const int source_lane = (lane & 15) * 2; + const unsigned rounded_bits = __half_as_ushort(rounded); + const half gate = __ushort_as_half(static_cast( + __shfl_sync(0xffffffffu, rounded_bits, source_lane))); + const half up = __ushort_as_half(static_cast( + __shfl_sync(0xffffffffu, rounded_bits, source_lane + 1))); + if (lane < 16) { + const float gate_f = __half2float(gate); + const half silu = __float2half(gate_f / (1.0f + expf(-gate_f))); + output[static_cast(route) * (n / 2) + tile * 16 + lane] = + __hmul(silu, up); + } + } else { + output[static_cast(route) * n + tile * 32 + lane] = rounded; + } + } +} + +// Qwen3.8 TP4 has ten K160 -> N2560 W2 routes. Keeping one route per warp +// retains split-K=1 accumulation, while grouping all routes for one N32 tile +// lets the CTA reduce them directly. Each route is rounded through FP16 before +// weighting, matching the former W2-output plus Triton-reduce path bit for bit. +__global__ void nvfp4_qwen38_w2_direct_reduce_kernel( + const half* __restrict__ input, const uint32_t* __restrict__ weights, + const half* __restrict__ scales, const int32_t* __restrict__ expert_ids, + const float* __restrict__ topk_weights, half* __restrict__ output) { + constexpr int kRoutes = 10; + constexpr int kK = 160; + constexpr int kN = 2560; + constexpr int kExperts = 512; + __shared__ half route_outputs[kRoutes][32]; + + const int lane = threadIdx.x & 31; + const int route = threadIdx.x >> 5; + const int tile = blockIdx.x; + const int expert = __ldg(expert_ids + route); + float accum[8] = {}; + + if (expert >= 0 && expert < kExperts) { + const int quadpair = (lane >> 2) & 3; + const int a_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int packed_col = + ((lane >> 2) & 3) * 8 + (lane & 3) + ((lane & 16) ? 4 : 0); + constexpr int kGroupsK16 = kK >> 4; + constexpr int kGroupsK8 = kK >> 3; + constexpr int kTilesN32 = kN >> 5; + constexpr size_t kWordsPerExpert = static_cast(kK) * kN / 8; + constexpr size_t kScalesPerExpert = static_cast(kK >> 4) * kN; + const uint32_t* expert_weights = + weights + static_cast(expert) * kWordsPerExpert; + const half* expert_scales = + scales + static_cast(expert) * kScalesPerExpert; + const half* input_row = input + static_cast(route) * kK; + +#pragma unroll + for (int group = 0; group < kGroupsK16; ++group) { + const size_t tile_group_base = + (static_cast(tile) * kGroupsK8 + group * 2) * 32 + packed_col; + const unsigned packed0 = __ldcs(expert_weights + tile_group_base); + const unsigned packed1 = __ldcs(expert_weights + tile_group_base + 32); + const size_t scale_index = + (static_cast(group) * kTilesN32 + tile) * 32 + packed_col; + const half scalar = __ldg(expert_scales + scale_index); + const half2 scale = + __hmul2(__halves2half2(scalar, scalar), __float2half2_rn(16384.0f)); + half2 decoded[8]; + dequant_e2m1x8(packed0, scale, decoded); + dequant_e2m1x8(packed1, scale, decoded + 4); + const unsigned* b = reinterpret_cast(decoded); + + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (a_row == 0) { + input01 = *reinterpret_cast(input_row + group * 16); + input23 = *reinterpret_cast(input_row + group * 16 + 8); + } + const unsigned* a0 = reinterpret_cast(&input01); + const unsigned* a1 = reinterpret_cast(&input23); + VLLM_SM70_MMA_8N8K4(accum, a0[0], a0[1], b[0], b[1]); + VLLM_SM70_MMA_8N8K4(accum, a0[2], a0[3], b[2], b[3]); + VLLM_SM70_MMA_8N8K4(accum, a1[0], a1[1], b[4], b[5]); + VLLM_SM70_MMA_8N8K4(accum, a1[2], a1[3], b[6], b[7]); + } + + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int index = pair * 4 + offset; + const int local_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + route_outputs[route][quadpair * 8 + local_col] = + __float2half(accum[index]); + } + } + } + } else if (lane < 4) { +#pragma unroll + for (int offset = 0; offset < 8; ++offset) { + route_outputs[route][lane * 8 + offset] = __float2half(0.0f); + } + } + __syncthreads(); + + if (route == 0) { + float weighted = 0.0f; +#pragma unroll + for (int selected = 0; selected < kRoutes; ++selected) { + weighted = fmaf(__ldg(topk_weights + selected), + __half2float(route_outputs[selected][lane]), weighted); + } + output[tile * 32 + lane] = __float2half(weighted); } } @@ -274,6 +443,25 @@ void launch_nvfp4_qpn_m1(torch::Tensor out, torch::Tensor input, reinterpret_cast(out.data_ptr()), n, k, broadcast_input); } +void launch_nvfp4_qwen38_w13_fused_swiglu(torch::Tensor out, + torch::Tensor input, + torch::Tensor weights, + torch::Tensor scales, + torch::Tensor expert_ids) { + constexpr int kN = 320; + constexpr int kK = 2560; + constexpr int kSplitK = 16; + const int routes = static_cast(expert_ids.numel()); + nvfp4_qpn_m1_sm70_kernel + <<>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weights.data_ptr()), + reinterpret_cast(scales.data_ptr()), + expert_ids.data_ptr(), + reinterpret_cast(out.data_ptr()), kN, kK, true); +} + void dispatch_nvfp4_qpn_m1(torch::Tensor out, torch::Tensor input, torch::Tensor weights, torch::Tensor scales, torch::Tensor expert_ids, bool broadcast_input, @@ -406,6 +594,112 @@ void nvfp4_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor input, C10_CUDA_KERNEL_LAUNCH_CHECK(); } +void nvfp4_qwen38_w2_direct_reduce_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, + torch::Tensor scales, + torch::Tensor expert_ids, + torch::Tensor topk_weights) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && weights.is_cuda() && + scales.is_cuda() && expert_ids.is_cuda() && + topk_weights.is_cuda(), + "nvfp4_qwen38_w2_direct_reduce_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + weights.scalar_type() == torch::kInt32 && + scales.scalar_type() == torch::kFloat16 && + expert_ids.scalar_type() == torch::kInt32 && + topk_weights.scalar_type() == torch::kFloat32, + "nvfp4_qwen38_w2_direct_reduce_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + weights.is_contiguous() && scales.is_contiguous() && + expert_ids.is_contiguous() && topk_weights.is_contiguous(), + "nvfp4_qwen38_w2_direct_reduce_out: tensors must be contiguous"); + TORCH_CHECK(out.sizes() == torch::IntArrayRef({1, 2560}) && + input.sizes() == torch::IntArrayRef({10, 160}) && + weights.sizes() == torch::IntArrayRef({512, 160, 320}) && + scales.sizes() == torch::IntArrayRef({512, 10, 2560}) && + expert_ids.numel() == 10 && topk_weights.numel() == 10, + "nvfp4_qwen38_w2_direct_reduce_out: shape mismatch"); + TORCH_CHECK(input.get_device() == out.get_device() && + input.get_device() == weights.get_device() && + input.get_device() == scales.get_device() && + input.get_device() == expert_ids.get_device() && + input.get_device() == topk_weights.get_device(), + "nvfp4_qwen38_w2_direct_reduce_out: device mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + nvfp4_qwen38_w2_direct_reduce_kernel<<<80, 320, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weights.data_ptr()), + reinterpret_cast(scales.data_ptr()), + expert_ids.data_ptr(), topk_weights.data_ptr(), + reinterpret_cast(out.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void nvfp4_qwen38_w13_fused_swiglu_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weights, + torch::Tensor scales, + torch::Tensor expert_ids) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && weights.is_cuda() && + scales.is_cuda() && expert_ids.is_cuda(), + "nvfp4_qwen38_w13_fused_swiglu_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + weights.scalar_type() == torch::kInt32 && + scales.scalar_type() == torch::kFloat16 && + expert_ids.scalar_type() == torch::kInt32, + "nvfp4_qwen38_w13_fused_swiglu_out: dtype mismatch"); + TORCH_CHECK(out.is_contiguous() && input.is_contiguous() && + weights.is_contiguous() && scales.is_contiguous() && + expert_ids.is_contiguous(), + "nvfp4_qwen38_w13_fused_swiglu_out: tensors must be contiguous"); + TORCH_CHECK(out.sizes() == torch::IntArrayRef({10, 160}) && + input.sizes() == torch::IntArrayRef({1, 2560}) && + weights.sizes() == torch::IntArrayRef({512, 2560, 40}) && + scales.sizes() == torch::IntArrayRef({512, 160, 320}) && + expert_ids.numel() == 10, + "nvfp4_qwen38_w13_fused_swiglu_out: shape mismatch"); + TORCH_CHECK(input.get_device() == out.get_device() && + input.get_device() == weights.get_device() && + input.get_device() == scales.get_device() && + input.get_device() == expert_ids.get_device(), + "nvfp4_qwen38_w13_fused_swiglu_out: device mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + launch_nvfp4_qwen38_w13_fused_swiglu(out, input, weights, scales, expert_ids); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void qwen38_shared_gate_exact_out(torch::Tensor out, torch::Tensor input, + torch::Tensor weight) { + TORCH_CHECK(out.is_cuda() && input.is_cuda() && weight.is_cuda(), + "qwen38_shared_gate_exact_out: tensors must be CUDA"); + TORCH_CHECK(out.scalar_type() == torch::kFloat16 && + input.scalar_type() == torch::kFloat16 && + weight.scalar_type() == torch::kFloat16, + "qwen38_shared_gate_exact_out: tensors must be float16"); + TORCH_CHECK( + out.is_contiguous() && input.is_contiguous() && weight.is_contiguous(), + "qwen38_shared_gate_exact_out: tensors must be contiguous"); + TORCH_CHECK(out.sizes() == torch::IntArrayRef({1, 2560}) && + input.sizes() == torch::IntArrayRef({1, 2560}) && + weight.sizes() == torch::IntArrayRef({1, 2560}), + "qwen38_shared_gate_exact_out: expected M1/N1/K2560 tensors"); + TORCH_CHECK(out.get_device() == input.get_device() && + out.get_device() == weight.get_device(), + "qwen38_shared_gate_exact_out: device mismatch"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + qwen38_shared_gate_exact_kernel<<<1, kQwen38SharedGateThreads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(input.data_ptr()), + reinterpret_cast(weight.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + void nvfp4_moe_qpn_mtp5_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor weights, torch::Tensor scales, torch::Tensor expert_ids, bool broadcast_input, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 8721632102..e46afd5851 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -534,6 +534,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor _gate_weight) -> ()"); ops.impl("sm70_f16_gate_mul_out", torch::kCUDA, &sm70_f16_gate_mul_out); + ops.def( + "qwen38_shared_gate_exact_out(Tensor(a!) out, Tensor input, " + "Tensor weight) -> ()"); + ops.impl("qwen38_shared_gate_exact_out", torch::kCUDA, + &qwen38_shared_gate_exact_out); + ops.def("sm70_gemm_import_cache(Tensor device_hint, str path) -> int"); ops.impl("sm70_gemm_import_cache", torch::kCUDA, &sm70_gemm_import_cache); @@ -697,6 +703,20 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("nvfp4_moe_qpn_m1_sm70_out", torch::kCUDA, &nvfp4_moe_qpn_m1_sm70_out); + ops.def( + "nvfp4_qwen38_w2_direct_reduce_out(" + "Tensor(a!) out, Tensor input, Tensor weights, Tensor scales, " + "Tensor expert_ids, Tensor topk_weights) -> ()"); + ops.impl("nvfp4_qwen38_w2_direct_reduce_out", torch::kCUDA, + &nvfp4_qwen38_w2_direct_reduce_out); + + ops.def( + "nvfp4_qwen38_w13_fused_swiglu_out(" + "Tensor(a!) out, Tensor input, Tensor weights, Tensor scales, " + "Tensor expert_ids) -> ()"); + ops.impl("nvfp4_qwen38_w13_fused_swiglu_out", torch::kCUDA, + &nvfp4_qwen38_w13_fused_swiglu_out); + // Keep the five-row verifier on a distinct schema so an old extension that // only supports the ten-route M=1 contract cannot be selected accidentally. ops.def( @@ -903,6 +923,25 @@ 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_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) -> ()"); diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index df8b0554d2..f4c7d77c41 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -816,3 +816,278 @@ Primary evidence is: - `.artifacts/qwen38_exact_decode80/gemv_router_ba_index_hc_bk256_qsa4_gdndual_gate_overlap_full/gemv_router_ba_index_hc_bk256_qsa4_gdndual_gate_overlap_full_i8192_o512_r5.json` - `.artifacts/qwen38_exact_decode80/gsm16_exact_decode80_official_xhigh/audit.json` - `.artifacts/qwen38_exact_decode80/gsm16_exact_decode80_official_xhigh/health.json` + +## 2026-09-04 current-main single-request decode trace + +The unified dual-compile/hybrid-PLE service was reprofiled at public-main SHA +`05910abb97446128a259fbd5fbe2bf9ece70a492`. The locked route is TP4/V2, +no MTP, FP16 activation/KV, checkpoint-native NVFP4 experts, full decode CUDA +Graph, prefix caching off, and input 8,192/output 513. One model load ran a +513-token low-overhead baseline outside the profiler capture and then captured +only a 32-token graph-node diagnostic. + +The 8K baseline measured `83.3749 tok/s`, or `11.9940 ms/token`; the accepted +short-prompt service point remains `86.07 tok/s`, so context length must stay +in every decode comparison. The node trace measured a `12.783 ms` middle-token +replay interval, `12.762 ms` GPU activity envelope, and only `0.050 ms` mean TP +replay-start skew. It covered `97.09%` of graph-node kernels and contained about +1,644 kernels/rank/token. Half of those kernels were shorter than 5 us. The +unprofiled decode samples reported 100% GPU utilization but only 140-150 W per +board, consistent with HBM traffic and small-kernel/graph-node issue cost rather +than FP16 compute saturation. + +Rank-average service attribution, which is not additive wall time because the +shared-expert stream overlaps the main stream, is: + +| Subsystem | Service ms/rank/token | +| --- | ---: | +| HyperConnection | 2.603 | +| NVFP4 MoE expert/router/activation | 2.203 | +| checkpoint-FP16 row GEMV | 1.441 | +| QSA sparse attention | 1.263 | +| remaining dense/cuBLAS, chiefly shared expert | 1.256 | +| fused GDN input | 1.064 | +| elementwise/metadata/copy | 1.022 | +| TP communication | 0.854 | +| GDN recurrent/core | 0.658 | +| LM head/sample | 0.582 | + +The checkpoint-FP16 HC down/up, fused GDN input, and remaining row-GEMV kernels +read at least 2.788 GB of weights per rank and token. Exact tensor sizes and +trace duration imply 552-596 GB/s for HC, 529 GB/s for row GEMV, and 714 GB/s +for fused GDN input. These are traffic lower bounds rather than NCU counters; +the current host blocks performance counters with `ERR_NVGPUCTRPERM`. The GDN +input is therefore not the first target. The ordered implementation candidates +are exact router projection/top-k, NVFP4 W2 plus weighted-reduce fusion, QSA +decode fusion, critical-path graph-node reduction around HC, and an exact or +guarded greedy LM-head route. Full-model startup is deferred until standalone +real-weight candidates project at least 0.4 ms/token combined savings. + +Raw reports remain outside Git under +`.artifacts/qwen38_nomtp_token_trace/`, including the `.nsys-rep`, exported +SQLite database, parsed per-token JSON/CSV/Markdown, route contract, and GPU +samples. + +### Exact post-trace candidates + +Three lossless single-token changes have passed focused operator gates after +the trace. They retain checkpoint FP16 activations and HC weights, native +NVFP4 expert weights, FP32 accumulation, and the existing FP16 materialization +boundaries: + +- The Qwen3.8 W2 kernel now forms each route's FP16 result before applying the + top-k weight and rank-ordered reduction in the same launch. Its real-weight + CUDA Graph gate is bitwise and projects `0.098 ms/token` savings over 48 + layers. +- HC up reuses the same 320-element low-rank vector across four independent + output rows. The selected row-four schedule is bitwise in all 128 changing + input cases and reduces the 96-call HC cycle from `2.139 ms` to `2.062 ms`, + saving `0.077 ms/token`. +- Qwen3.8 M1 `all_reduce_sum2` now reuses the registered SM70 TP4 push buffers. + Both paths first form the local FP16 sum, accumulate ranks 0 through 3 in + FP32, and round once to FP16. The four-rank CUDA Graph gate is bitwise for + integer, model-distribution, and signed-zero patterns. Forty-eight + collectives fall from `0.459 ms` to `0.136 ms`, saving `0.323 ms/token`. + `VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1=0` is the rollback. +- Single-row QSA decode now performs one coarse score-radix pass, compacts only + that bucket, and refines the remaining radix bytes in shared memory. The + final increasing-index scan is unchanged, so lower-index score ties and the + downstream accumulation order remain exact. Twelve real-shape launches at + lengths 2,048-2,169 fall from `0.2169 ms` to `0.1238 ms`, saving + `0.0931 ms/token` or 1.75x. Random scores, dense ties, signed zero, Inf/NaN, + the 2,304-entry boundary, and the 2,305/4,096/16,384 device fallback are + bitwise equal to the original selector. Multi-row prefill retains the + original kernel. Source-overlay validation may supply the same compiled + fragment through `VLLM_SM70_QSA_TOPK_LIBRARY`; release wheels link it into + `_C_stable_libtorch` normally. + +The isolated savings sum to `0.591 ms/token`; they are not an end-to-end TPOT +claim because the shared-expert and main streams overlap. One full-model A/B is +still required before treating the operator gains as service throughput. + +Privileged NCU counters confirm why HC needs traffic/issue improvements rather +than lower precision. The down projection reaches `488 GB/s` DRAM throughput +with `24.23%` achieved occupancy and spends `86.11%` of scheduler cycles with +no eligible warp. HC up reaches `481 GB/s`, `64.31%` occupancy, and `63.45%` +no-eligible cycles. More warps, split-K down, down row tiling, and the SGLang +persistent atomic-grid HC implementation are slower on V100. HC +combine-plus-RMSNorm is already only about `0.305 ms` per 96-call graph cycle; +an 8-warp variant saves just `0.014 ms` and changes the reduction result, so it +is rejected. A warp-per-dot HC-up kernel is `0.132 ms/token` slower and changes +89 of 245,760 FP16 outputs by at most `0.000488`. A same-precision FP16 Tensor +Core/QPN layout is also `0.012-0.035 ms/token` slower, consumes about 0.6 GiB +more packed weights per rank, and does not reproduce the established FP16 +materialization boundary. Both are rejected. The retained HC changes do not +quantize FP16 tensors or relax any quality gate. + +Further exact HC screens close the inexpensive schedule space. Bypassing L1 for +streaming weights is bitwise but `0.054 ms/token` slower. Changing Triton +pipeline stages is bitwise and neutral within `0.002 ms/token`; 8/16-row HC-up +tiles and paired-stream prefetch are bitwise but `0.043-0.097 ms/token` slower. +Larger down reduction tiles save at most `0.022 ms/token` while changing FP16 +outputs by one ULP, so they are rejected. Fusing the attention output projection +with HC combine, followed by an exact norm-only kernel, is bitwise for both +multi-stream and normalized outputs but is `0.013 ms/token` slower over 48 +calls. These paths should not be rescanned without a different kernel +architecture. + +### Exact TP4 HyperConnection compute sharding + +The next retained HC candidate changes work placement, not model precision. +For each M=1 HC down projection, rank `r` computes low-rank rows +`[80r, 80(r+1))` and injection row `320+r` directly from the existing +replicated checkpoint-FP16 weight. A rank-ordered push all-gather reconstructs +the original 320 low-rank and four injection values. Each rank then computes +the corresponding 2,560 rows of the FP16 HC-up projection, and a second push +kernel applies the established FP16 gate boundary, FP32 sigmoid and +rank-ordered FMA, and final FP16 materialization. The implementation keeps the +full weights resident, so prefill and unsupported cases use the original +replicated path without a weight-loader or memory-layout change. + +On four V100-SXM2-32GB GPUs, the real-shape 96-HC CUDA Graph cycle falls from +`2.042378 ms` to `1.748982 ms`, saving `0.293396 ms/token` or 16.78%. All +block and injection outputs are bitwise equal on all four ranks. A separate +production-dispatch smoke covers 16 changing inputs through the registered +custom op and CUDA Graph lifecycle; all four ranks report zero FP16 bit +mismatches. The route requires the existing checkpoint-FP16 HC opt-in, exact +Qwen3.8 topology, fully connected TP4 SM70 custom all-reduce, and registered +push buffers. Otherwise it falls back before launching a sharded kernel. + +This candidate does not use FP8, INT8, QPN, altered activation types, or a +reduced-precision accumulator. Together with the preceding isolated exact +screens, projected operator savings are `0.884 ms/token`; this is still not an +end-to-end throughput claim, and it does not by itself establish the 100 +tok/s target. + +Two additional no-lower-precision screens were rejected. A deterministic +E512/K10 top-10 selector preserves all outputs bitwise across random inputs, +dense ties, signed zero, NaN, and infinities, but loses `0.023 ms/token` with +hot logits and `0.049 ms/token` after a 64-MiB L2 scrub. GDN input row tiling +is bitwise but saves only `0.0046 ms/token`. Checkpoint-native NVFP4 W13 +split-16 retains FP32 MMA accumulation and FP16 output but changes FP32 +summation grouping; it differs from split-8 by one FP16 ULP in about 0.28% of +sampled outputs, so it is not enabled without a full model quality gate. + +### Hidden-coordinate HC sharding, 2026-09-05 + +The next exact M=1 candidate assigns each TP rank 640 hidden coordinates and +computes all four branch gates for those coordinates. It preserves the +checkpoint FP16 weight layout, the two-K-warp FP32 reduction, the FP16 gate +boundary, FP32 sigmoid and branch-ordered FMA, and the final FP16 output. +The following collective gathers final hidden slices instead of branch gates: +each rank sends 1,280 rather than 5,120 bytes to each peer. No extra weight +copy or precision change is introduced, and prefill is unchanged. + +It uses the existing `VLLM_SM70_QWEN38_FUSED_HC_FP16` opt-in and exact TP4 +admission checks. A source-matched custom-AR extension enables the new +`sm70_qwen38_hc_output_allgather` op; an older extension retains the existing +gate-sharded path. Capability discovery and dispatch use the DSO that owns the +opaque communicator, never a different extension's fallback symbol. The new +gather reuses the isolated HC channel, not the concurrent MoE channel. + +The initial screen loads all 96 real HC weight pairs, not a repeated layer-0 +weight. Four V100-SXM2-32GB ranks each report zero FP16 bit mismatches for +block and injection outputs over 16 changing inputs. After 1,000 warmup graph +replays, three alternating paired timing groups give the following medians: + +| Variant | 96 Mix calls (ms) | Change from control (ms) | +| --- | ---: | ---: | +| Current gate-sharded control | 1.743988 | — | +| Hidden shard, two hidden rows / eight warps | 1.703158 | -0.040830 | +| Producer-only down publication, coalesced revision | 2.037357 | +0.293369 | +| Exact down partials + fused tail/gather, one part | 1.842709 | +0.098720 | +| Same, two parts | 1.850873 | +0.106885 | +| Same, four parts | 1.864315 | +0.120327 | + +Only hidden sharding is retained. The first producer-only version was still +slower at 2.229951 ms. Coalescing its peer writes reduced that overhead but +did not beat the control. Fixed-order down splitting also remained slower +after half2 loads and a one-warp gather tail. None of those losing prototypes +is part of the production dispatch. Alternative hidden tiles / warp counts +were bitwise but slower than the selected two-row/eight-warp schedule. + +These are Mix-only graph measurements: they exclude HC combine/RMSNorm and +must not be subtracted directly from the 2.658-ms full-HC trace service sum. +The initial prototype improvement is about 2.3%, not the initial 20% screening +target and not an end-to-end tokens/s claim. + +The committed production implementation (`aaf63696b6`) subsequently passes +the registered-op gate: 96 real weight pairs x 16 changing inputs x four +ranks, including 512 graph replays overlapping the actual sum2 route on an +auxiliary stream. All HC block, injection, and sum2 outputs have zero FP16 bit +mismatches. The independent hidden-shard GPU unit test passes, as do all 13 +CPU dispatch tests. With Torch `2.10.0+cu128`, runtime CUDA `12.8`, and the +SM70 extension compiled by NVCC `12.0.140`, three paired timings are: + +| Production dispatch | Paired samples (ms) | Median (ms) | +| --- | --- | ---: | +| Existing gate shard | 1.738315 / 1.739291 / 1.738595 | 1.738595 | +| New hidden shard | 1.689020 / 1.690590 / 1.690003 | 1.690003 | + +The final Mix-only saving is **0.048592 ms (2.79%)**; each variant's sample +range is below 0.1%. No full-model startup is justified by this small increment +alone. Combine it with other admitted exact candidates for the next matched +endpoint gate; natural-output quality and the required 256K boundary remain +part of endpoint promotion. The 100-tok/s target remains unproven by this +change. + +Reproduce the production-dispatch gate without loading the entire model: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SM70_TP4_PUSH_ALLREDUCE=1 \ + .venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=4 \ + benchmarks/kernels/benchmark_sm70_hc_tp4.py \ + --model /path/to/Qwen3.8-Flash-Next-NVFP4 --out /path/to/hc-result.json +``` + +Use a source-matched wheel or set `VLLM_SM70_CUSTOM_AR_LIBRARY` to an extension +that contains both the new op and the complete communicator lifecycle. The +benchmark compares old and new registered-op dispatch, checks all 96 real +weight pairs, overlaps HC with the actual sum2 CUDA Graph route on an auxiliary +stream, and reports three paired Mix-only timings separately from correctness. + +### Exact fused HC up/mix/gather candidate (2026-09-05) + +The source now includes the selected 160-CTA FP16 up/mix/gather kernel. It +uses 128-bit weight/input reads, parallel branch sigmoids, and the same +eight-term FP32 FMA chains, XOR reduction tree, FP16 gate boundary, and +branch-ordered FP32 mixing. Each CTA publishes exact FP16 outputs together +with a generation tag; its two packet slots and generation counter are +isolated from both legacy HC collectives and the auxiliary MoE channel. +The existing communicator allocation grows by 21,120 bytes per rank; there +is no weight copy, additional communicator, or new user tuning switch. + +The existing `VLLM_SM70_QWEN38_FUSED_HC_FP16` opt-in and TP4/SM70/M=1 gates +still apply. A source-matched extension selects fused up/mix/gather; an older +extension retains hidden-sharded split up/gather, or the legacy gate-sharded +route if needed. Optional-op capability and dispatch must come from the DSO +that owns the communicator, including the extended allocation layout. + +The preceding **prototype** complete-HC screen measures `2.108826 -> +1.999374 ms` (5.19%) with bitwise intermediate/final outputs. The subsequent +registered production gate at `0303b82d1e` measures **`2.109529 -> 1.994807 ms` +(5.44%)**. All four ranks pass the 16-input intermediate/final checks, 512 +auxiliary sum2 replays, and post-timing checks after packet generation wrap. +Fused samples are `1.994807/1.993735/1.996370 ms`, versus split-hidden +`2.109556/2.109529/2.109242 ms`. Runtime is Torch `2.10.0+cu128`, CUDA `12.8`, +TP4 V100-SXM2-32GB; the sidecar was compiled with NVCC `12.0.140`. +The `1.5-ms` whole-HC target and endpoint speed are not established by this +microbenchmark. The old `2.658-ms` whole-model trace uses a different scope +and must not be compared directly to it. + +Run the complete registered-op gate, without loading attention/MoE weights: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 CUDA_DEVICE_ORDER=PCI_BUS_ID \ + VLLM_SM70_TP4_PUSH_ALLREDUCE=1 VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1=1 \ + .venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=4 \ + benchmarks/kernels/benchmark_sm70_hc_full_chain.py --fused-up \ + --model /path/to/Qwen3.8-Flash-Next-NVFP4 --out /path/to/hc-full-result.json +``` + +This compares forced split-hidden and fused registered dispatch, includes all +HC norms and final projections, checks 16 changing inputs and 512 auxiliary +sum2 graph replays, then rechecks outputs after timing crosses packet-tag +wrap. Timings exclude the auxiliary stress workload. Both this benchmark and +the older Mix-only gate explicitly freeze their control routes so a newer +extension cannot silently replace both sides of the comparison. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index ce84e94a98..bde32d70cc 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -44993,3 +44993,355 @@ Interpretation: 200. The service retained the 47.684-GiB file-backed PLE mapping with about 1.3 GiB worker RSS; startup's cgroup peak includes reclaimable/shared file mappings and did not trigger `systemd-oomd` after sequencing was repaired. + +## 2026-09-04 Qwen3.8 exact no-MTP decode continuation + +- The current TP4/V2/no-MTP baseline uses Qwen3.8-Flash-Next-NVFP4 on four + V100-SXM2-32GB GPUs, FP16 activation/KV, Flash-V100, full CUDA Graphs, + hybrid mmap-prefill plus pinned-decode PLE, 8,192 input tokens, and 513 + generated tokens. The accepted result is `85.136 token/s`, or + `11.746 ms/token`; its output token IDs exactly match the preceding control + (`7385dac...`). The active target remains `100 token/s` without changing + arithmetic precision or enabling MTP. +- A graph-node trace at source `234d6bea91` measures `12.559 ms/token` with + tracing overhead. The leading additive categories are row GEMV + `1.439 ms`, GDN input `1.059 ms`, W13 `0.656 ms`, HC local down + `0.564 ms`, LM head `0.552 ms`, router `0.525 ms`, HC local up + `0.492 ms`, W2 `0.480 ms`, QSA split-K `0.440 ms`, and HC combine/norm + `0.402 ms`. Shared-expert auxiliary GEMVs overlap routed MoE and are not + added to the critical path a second time. +- The production sidecar now includes the existing exact direct-W2 reduction + operator that the prior deployed binary lacked. Its production-shape graph + screen is bitwise equal and moves the W2 plus weighted-reduce chain from + `0.5062` to `0.4082 ms/token`, a projected `0.0980 ms/token` saving. This + projection is intentionally held for a combined model startup. +- Two narrow SM70 PLE M=1 kernels remove generic tensor plumbing without + changing data types or rounding boundaries. The exact ngram-2/3 ID kernel + passes 256/256 random and EOS-boundary comparisons and moves + `0.08169` to `0.00435 ms/token`. The depthwise dilated-convolution/state + kernel is bitwise for normal, no-initial-state, and graph-padding cases and + moves `0.04084` to `0.00543 ms/token` while retaining native `F.silu`. +- Fusing SiLU into Triton was rejected because the approximation changed two + FP16 values by up to `1.22e-4`; the admitted path keeps the original native + SiLU rounding. Cooperative HC combine/down, sum2/combine fusion, a + deterministic replacement router, SGLang's atomic persistent-HC design, + and fine-grained down-project/push pipelining were also rejected as slower, + nondeterministic, or deadlocking. They must not be retried without a new + schedule or arithmetic proof. +- The direct-W2 and PLE screens project about `0.211 ms/token` combined. No + full-model speed result is claimed yet: these changes are deliberately + batched with further exact hot-path work so model loading is not repeated + for a sub-millisecond projection. +- The official Qwen3.8 PLE gate and merged key/value projection were screened + separately on V100 before porting. The gate saved only `0.0022 ms` at the + production M=1 shape and changed thousands of FP16 elements. Merging the two + projections saved only `0.0031-0.0034 ms` and changed 11 of 12,800 FP16 + outputs. Both are rejected for the no-precision-loss lane. +- The checkpoint-native interleaved W13 layout places each gate/up pair in one + N32 CTA. The admitted decode epilogue retains the existing FP32 split-16 + accumulation, rounds each projection to FP16 at the same boundary, then + evaluates the existing `expf` SiLU and FP16 multiply using warp shuffles. + Across 48 production-shaped layers, three CUDA Graph runs are bitwise equal + and save `0.0656-0.0720 ms/token` (`0.6045-0.6106` to + `0.5386-0.5388 ms/token`). Both variants use 60 registers/thread and 2 KiB + shared memory. Older extensions without the new op fall back to the prior + exact two-kernel route. +- Direct-W2, exact PLE, and fused W13/SwiGLU now project about + `0.277-0.283 ms/token` combined. This remains a projection rather than an + end-to-end claim; retain it for the next material combined model startup. +- A shared-expert gate screen localizes the old M=1 path to separate scalar + projection/reduction, sigmoid, and 2,560-element output-multiply kernels. + The prior `sm70_f16_gate_mul_out` candidate is rejected for this lane + because it rounds products to FP16 and omits the eager FP16 linear boundary; + on 48 real layer-0-shaped cases it changed 13,947 output elements by up to + `0.001953125` despite improving `0.4278 -> 0.1231 ms/token`. +- The accepted shared-gate candidate retains FP32 FMA, explicitly rounds the + scalar linear output and sigmoid result to FP16 at the original boundaries, + and performs the final FP16 output multiplication in one CTA. Across all 48 + real checkpoint gate weights and 512 changing inputs, the rounded gate and + all 2,560 output elements are bitwise equal (`0` mismatches). The production + sidecar/facade route measures `0.4261 -> 0.1253 ms/token`, saving + `0.3009 ms/token` for the isolated 48-layer chain without changing weight, + activation, accumulation, or output precision. Because shared experts can + overlap routed MoE, this is not counted one-for-one as endpoint TPOT until a + combined full-model run measures the reduced contention. +- An exact HC-up projection/push experiment reproduced Triton's two-warp + K-split and XOR reduction tree, eliminating all 22 mismatches from the older + prototype on every TP4 rank. Expanding the fused collective from 32 to 80 + CTAs nevertheless regressed the 96-HC chain from `1.7440` to + `2.3930 ms/token`; P2P polling and CTA overhead dominate the saved launch. + The 80-CTA fusion is rejected and must not replace the current split path. +- vLLM PR 55309's QSA output-gate fusion was adapted to this tree's split-K, + E4M3-scale, and SM70 XQA branches. The generic split-merge and direct-write + paths preserve the compiled model's FP16/BF16 attention-output boundary, + then evaluate sigmoid and multiplication in FP32 before the final store. + The TP4 decode/split-64 and large-batch/split-1 tests are bitwise equal to + the previous separate compiled gate. A 12-QSA-layer CUDA Graph screen at + 8K context improves `0.35888 -> 0.34250 ms/token`, saving + `0.01639 ms/token` (`1.048x`) with zero differing elements. +- The same upstream PR's PLE outer-residual patch is not directly portable as + an additional decode optimization here. This tree already compiles the two + PLE residual additions into one three-input FP32 pointwise kernel. Its exact + M=1 short-convolution deliberately retains native `F.silu`: the earlier + Triton SiLU fusion changed FP16 results. Moving the add across the custom-op + boundary without also replacing native SiLU would not remove a launch, so + no PLE residual source change is admitted from this PR. + +## Exact HC three-direction screen, 2026-09-05 + +- Owner: `codex/v100-qwen38-nomtp-token-trace-20260903-173451`, public Draft + PR [#481](https://github.com/1CatAI/1Cat-vLLM/pull/481). Pre-change source + `30f81105621e9f39e6b3bf9d816f77d63acd8307`; current integration merge base + `fbcef6e2f959e95bbe4ca807931abfa2393546e7`. Work stays in the owned worktree; + no direct push to `main` and no change to another task's running API. +- Frozen operator contract: RadixArk/Qwen3.8-Flash-Next-NVFP4, all 48 layers' + attention and MLP HC pairs (96 distinct weights), TP4 V100-SXM2-32GB, M=1, + checkpoint FP16 weights/inputs/outputs, FP32 arithmetic, no MTP. These are + HC Mix-only CUDA Graph cycles, not full HC, prefill, or endpoint TPOT. +- Implemented and screened all three research directions: hidden-coordinate + ownership, producer-only down publication, and exact logical-lane down + splitting with fused reduction/communication. All 96 pairs x 16 changing + inputs x four ranks are bitwise for block and injection outputs. +- Only hidden ownership is retained. It gathers 640 final FP16 values per + rank instead of 2,560 gates, with no additional resident weight copy. Two + stable paired screens show `1.745654 -> 1.702919` and + `1.743988 -> 1.703158 ms` per 96 Mix calls, saving `0.041-0.043 ms` (about + 2.3-2.4%). Three paired groups use 150 replays each after 1,000 warmups; + the second screen's range is below 0.2% for each retained variant. +- Rejected: direct per-row publication (`2.229951 ms`), its coalesced revision + (`2.037357 ms`), and exact one/two/four-part down with the improved + half2-load/one-warp gather tail (`1.842709/1.850873/1.864315 ms`). These are + slower than the `1.743988-ms` matched control despite preserving precision. + The coalesced publication and tail revision were targeted responses to the + first screen, not new full-model startups. Do not rescan them unchanged. +- Hidden two-row/four-warp, four-row/eight-warp, and four-row/sixteen-warp + schedules are also bitwise but slower at `1.738779/1.741660/1.719310 ms`. + Select two hidden rows and eight warps; do not conflate row tiling with a + change to the K-reduction tree. +- A separate publication prototype single-GPU four-peer emulation passed + 18 real-weight cases including generation 65535/65536 and signed 32-bit + wrap. This is arithmetic/protocol evidence only, not proof of distributed + speed or grounds to retain a slower publisher. +- Source integration retains the existing HC opt-in, exact shape gate, and + older-extension fallback. A new capability check follows the communicator's + owning DSO, preventing an old sidecar from borrowing a new base-wheel op. + The four ownership/capability cases pass; the complete focused CPU dispatch + suite is `13 passed`. +- Reproducible production gate: + `benchmarks/kernels/benchmark_sm70_hc_tp4.py --model MODEL --out RESULT`, + launched with four torchrun ranks and the source-matched extension. It uses + the registered HC custom op, compares forced old dispatch with new dispatch, + and checks concurrent sum2 on an auxiliary stream. Detailed launch examples + and measurement limitations are in + [the decode guide](sm70_qwen38_nvfp4_decode.md#hidden-coordinate-hc-sharding-2026-09-05). +- Local raw evidence: `.artifacts/hc_hidden_shard/stages_result.json`, + `coalesced_result.json`, `single.log`, `dispatch_test.log`, and + `build_production.log`. The initial `result.json` had large idle-clock jitter + and is not accepted timing evidence. Warmup was added before the stable + screens. No full-model load has been performed for this HC screen. +- The first production validation attempt was stopped by its owner before + timing when another task began a TP4 model run on GPUs 0-3. Its partial log + is `production_run.log`, not a failed numerical gate or performance result. + The guarded runner now checks both locks and actual device memory before + launch. Subsequent exit-75 lock waits are not model/GPU test attempts. +- Final production gate at source `aaf63696b6`: all 96 real weight pairs x 16 + changing inputs x four ranks pass bitwise, including 512 graph replays with + the actual sum2 route on an auxiliary stream. HC block, injection, and sum2 + each have zero differing FP16 bits on every rank. The independent GPU + hidden-shard test is `1 passed, 18 deselected`; CPU dispatch remains + `13 passed`. +- Three production paired samples are control + `1.738315/1.739291/1.738595 ms` and hidden + `1.689020/1.690590/1.690003 ms`. Medians are + `1.738595 -> 1.690003 ms`, saving **0.048592 ms (2.79%)** per 96 Mix calls; + ranges are below 0.1%. Runtime is Torch `2.10.0+cu128`, CUDA `12.8`, with + the SM70 sidecar compiled by NVCC `12.0.140`. Binary SHA256: + `a1fa27c23aea3ee2a7030017ee404c9d2bcb1f3c03889461a070c1f4daded4dd`. + Results/logs: `.artifacts/hc_hidden_shard/production_result.json` and + `production_final.log`. Result SHA256: + `6bf2c047430e586bf1814fbf8ae0fd09a335d4c59decc8d6fff4c5ca6aa37750`. +- All task-owned GPU tests and lock holders exited after validation. Other + tasks' model workers/API were not stopped. No full-model startup or endpoint + was launched for this small HC increment. The next combined full-model + quality/performance gate remains pending; do not claim 100 tok/s or promote + an endpoint from the isolated 0.049-ms saving. + +## Full HC <= 1.5 ms target, 2026-09-05 + +- The user explicitly set the next target to full HyperConnection latency + below `1.5 ms/token`, with no precision reduction. Preserve checkpoint FP16 + weights/activations, FP32 accumulation, and the established rounding/order + contract. A Mix-only result does not satisfy this target. Final acceptance + requires matched full-model trace attribution and output quality, not just + an isolated graph score. +- The old `2.658-ms` trace bucket was grouped by HC kernel names. It excluded + the final mixer's ordinary down/up projections (classified as dense work). + The new semantic HC microbenchmark includes these as well: 96 layer Mix + pairs, 95 combine/norm calls, two grouped input norms, the PLE boundary's + separate combine, and the final projection/SiLU/gate-mix path. Attention, + MoE, and PLE computation are excluded; attention/MoE outputs are fixed + external inputs. This is not a full-model run or an endpoint TPOT metric. +- Frozen source `50f9fbe3749ecd673fee1aef361afd8db98e914a`, same Torch + `2.10.0+cu128`, CUDA runtime `12.8`, TP4 V100-SXM2-32GB, and source-matched + HC sidecar as the preceding screen. Public integration advanced to + `9ed8697ac0` during this work; the candidate kernel source was not changed + to mix unrelated integration changes into the comparison. +- Complete HC microbenchmark medians are gate-sharded `2.277335 ms` and + current hidden-sharded `2.106689 ms`. Hidden samples are + `2.103712/2.106914/2.106689 ms`. All intermediate state, normalized state, + block input, injection, and final-mixer output tensors are bitwise across + all four ranks and 16 changing input cases. This cannot be reported as + `2.658 -> 2.107 ms` improvement: source scope, tracing, and workload differ. +- Independent component graph medians are down `0.486018 ms`, down gather + `0.305125 ms`, hidden up/mix `0.505760 ms`, output gather `0.212166 ms`, + 95 combine/norm calls `0.311712 ms`, and final mixer including its norm + `0.032160 ms`. Do not add these to close the complete graph: dependencies, + cache state, and the final norm overlap between component scopes differ. +- `benchmarks/kernels/benchmark_sm70_hc_full_chain.py` provides the portable + complete-HC registered-op gate. It initializes cuBLAS before capture and + checks exclusive GPU process ownership around timing groups. The initial + artifact harness missed cuBLAS warmup and failed at handle creation during + capture; this was corrected once before the accepted baseline. No model + startup was involved. +- Exact physical down expansion (128/256 threads with original 40-term FMA + chains and XOR tail tree) passes the full four-rank bitwise gate. Another + task entered the GPUs during timing; its large timing variance is not + accepted performance evidence. Neither CUDA down variant is admitted. + Ownership checks now also run during timing, not just before launch. +- A materially different combine/norm + down prototype preserves Triton's + original two-axis norm reduction and uses four producer CTAs with + release/acquire readiness, instead of the old changed-order cooperative + reduction/global grid barrier. Cooperative launch bounds the residency + requirement. It passes the full bitwise gate but loses: matched hidden + `2.105330 ms`, fused without prefetch `2.134842 ms`, fused with four-chunk + prefetch `2.125169 ms`. Do not admit these variants. +- One targeted follow-up prefetches all 40 immutable weight chunks before + consuming normalized input. A tensor-gather implementation was rejected at + compile/resource inspection (32 KiB dynamic shared memory, 255 registers, + 480-byte stack frame, large generated code) without a GPU timing trial. + The static-register revision compiles with 125 registers, 64 bytes dynamic + shared memory and no stack/local spill. Its full-chain gate passes bitwise, + but matched medians are `2.109467 -> 2.159213 ms` (regression). Stop this + combine/norm + down fusion direction; do not repeat the failed variants. +- A separate exact FP16 layout prototype packs four successive down chunks + into each 16-byte vector read and interleaves up's four branch rows by + hidden coordinate. Arithmetic order is unchanged; decode-only packed + shards cost `316538880 bytes` (`301.875 MiB`) extra per rank if both are + retained. The packed down compiles with 31 registers, 16 bytes shared + memory, no stack/local spill. All full-chain variants pass bitwise. Matched + medians are hidden `2.114089 ms`, packed down `2.174867 ms`, packed up + `2.100613 ms`, both `2.159398 ms`. Down/both are rejected; up's isolated + `0.013476-ms` saving is too small to justify admission on this evidence + alone. No production weight-loader or kernel route has changed for it. +- The portable registered-op benchmark also passes 16 changing inputs on all + four ranks. Medians are gate `2.277540 ms`, hidden `2.111058 ms`; hidden + samples are `2.109891/2.111399/2.111058 ms`. It agrees with the artifact + harness, and remains an isolated complete-HC measurement, not endpoint TPOT. +- Local evidence lives under `.artifacts/hc_full_chain/`: + `baseline.json`, `baseline_warm.log`, `down_schedules.json` (contended + timing, quality only), `fused_norm_down.json`, `fused40.json`, `packed.json`, + `public_baseline.json`, compiler logs and resource artifacts. The guarded + queue completed all three pending jobs and released its GPUs. Other + API/model tasks were not terminated. No full-model startup was involved. + The `1.5 ms` goal remains active and unachieved. Next screen targets hidden + up/local-mix/output-gather fusion with private per-CTA communication epochs, + distinct from the previously rejected branch-sharded/global-counter fusion. +- The new hidden-sharded 80-CTA CUDA up/mix/output-gather prototype uses exact + FP16-value-plus-generation packets and independent two-slot CTA epochs; + no global completion counter or sentinel-value substitution. Its full-chain + 16-input/four-rank gate passes bitwise. Matched medians are control + `2.105945 ms`, CUDA up/mix with separate gather `2.265607 ms`, fused gather + `2.180970 ms`. Fusion saves `0.084637 ms` relative to the CUDA split version, + but the projection schedule loses more: the net result is slower and is + rejected. Evidence: `up_gather.json`, source snapshot `up_gather80.cu`. + A bounded 160/320-CTA follow-up tests whether more projection parallelism can + retain the communication saving; each tile has its own private peer buffer + and generation counters. This is not a new production route. +- The bounded scalar-load follow-up is also not admitted: control + `2.108150 ms`; 160-CTA local/fused `2.161794/2.103446 ms`; 320-CTA + local/fused `2.139696/2.183004 ms`. The best saving is only `0.004704 ms`. + All four ranks pass the 16-input bitwise gate and a second comparison after + generation `146593` (two 16-bit wraps). Evidence: `up_gather_tiled.json`. +- SASS inspection identifies scalar U16 weight loads and shared-lora staging + in the CUDA prototype. A separate LDG128 revision removes staging while + preserving all arithmetic/rounding. It passes the initial four-rank + 16-input gate, but another task enters during timing; the benchmark rejects + the sample and exits. This is a contention-aborted result, not a numerical + failure or a speed claim (`up_gather_vector.log`). A bounded follow-up also + distributes the four branch sigmoids over four times as many active lanes + at the existing gate-materialization barrier, without an extra barrier or + changed FP16 boundary. Its 80/160-CTA full-chain gate is pending. +- The repeated contention is localized to a separate GPU reservation held + across another suite's model restarts. The guarded runner now honors that + existing flock as well as this task's GPU locks, without truncating the + other lease file. Do not enter the reserved suite's between-model gaps, + interrupt it, or accept contended timing. The vector/parallel-gate kernel + compiles without spills; its queued GPU screen remains pending. +- The vector-load/parallel-gate screen subsequently completed under the + shared reservation. It produces the first material complete-chain win in + this follow-up: control `2.108826 ms`, 160-CTA local-only `2.068084 ms`, + 160-CTA fused `1.999374 ms`. The full-chain saving is **`0.109452 ms` + (`5.19%`)**; fused samples are `1.999995/1.999374/1.998002 ms`. The 80-CTA + fused version is `2.082618 ms` and is not selected. There is no packed-weight + copy. All intermediate/final outputs pass bitwise on four ranks over 16 + changing input cases and again after generation `146593` (two tag wraps). + Evidence: `up_gather_vector.json` and `up_gather_vector_final.log`; source + SHA256 `10d65cb0b979a51b3e6cf712dd3c535d93f66e4b418909adeec1616077c4def5`. +- This is still an artifact prototype, not a registered production-path or + whole-model result. Next: port the selected 160-CTA kernel/private channel + with extension-capability fallback, then validate production dispatch and + actual auxiliary-stream sum2 coexistence. Batch the full-model trace and + output-quality gate with further material changes. Do not claim the old + full-model HC bucket moved `2.658 -> 1.999 ms`, or that the `1.5-ms` target + was achieved. All task-owned GPU tests/queues exited; other tasks continue. + +## Registered HC up/mix/gather port, 2026-09-05 + +- Previous goal turn made progress: the complete-HC prototype improved by + `0.109452 ms`, with bitwise evidence. The current goal remains full HC below + `1.5 ms/token` on matched whole-model trace, no MTP or precision reduction. +- Port only the selected 160-CTA/vector-load/parallel-gate kernel. Preserve + original FP32 FMA/reduction order and FP16 boundaries. Append a 21,120-byte + private packet/counter region to the existing communicator allocation; + legacy HC and auxiliary MoE layouts are unchanged. Start from zero counters + and publish generation one first. No extra weight copy or public switch. +- Add `sm70_qwen38_hc_up_mix_allgather` to production binding, owner-DSO + facade, communicator, and model dispatch. Keep split-hidden and gate-sharded + fallbacks for older owner DSOs. Never borrow the new op from another DSO, + which may have allocated a different buffer extent. +- Extend the complete-HC benchmark with `--fused-up`, forced control dispatch, + 16 changing-input checks, 512 replays with actual sum2 on an auxiliary + stream, and a post-timing bitwise check after generation wrap. The legacy + Mix-only benchmark also explicitly disables the new route in its controls. +- CPU dispatch/owner/fallback suite: **20 passed**. Ruff on affected Python + files and changed-line clang-format pass. Source-matched sidecar compiles; + dynamic `RankData` indexing initially introduced a 64-byte stack frame. + Constant parameter indices remove it before GPU testing: selected kernel + uses 31 registers, 192 bytes shared memory, zero stack or spill traffic. +- Evidence is under `.artifacts/hc_up_fused_production/`: `cpu_tests.log`, + `build_final.log`, source-only sidecar builder, guarded `run_when_idle.sh`. + Registered GPU gate is pending at this update; no whole model was started. + Public integration was fetched at `2b89b77e3882423d1c93e01faf8c1db43f6650f4`; + keep the candidate's existing integration base `fbcef6e2f9` frozen for this + paired screen rather than mixing unrelated model-route updates into it. +- Registered gate completed at `0303b82d1ec8fd9549d75018995939bbee63846e`: + full semantic HC split `2.109529 ms` -> fused **`1.994807 ms`**, saving + **`0.114722 ms` (`5.44%`)**. Split samples are + `2.109556/2.109529/2.109242`; fused `1.994807/1.993735/1.996370 ms`. + All four ranks have zero HC/intermediate/final and sum2 bit mismatches over + 16 changing input cases, 512 auxiliary sum2 replays, and post-timing checks + after generation wrap. `result.json` SHA256: + `b9524acfe04ea92ca3836a404ae590284dc6ea0b8ee4ddfd3a3488e9653a9996`. + Production binary SHA256: + `5b1ee678bebf6a8fcdb008d5832cfd8ca3d6978558291ec9fe54ec2b9f6cf1bf`. + All test processes exited; no full-model startup. This is a production-op + microbenchmark, not a whole-model trace or endpoint acceptance. +- Next bounded screen is an exact down/gather packet fusion, under + `.artifacts/hc_down_packet/`. It is materially different from the old + rejected 80-CTA/16-byte-sentinel fusion: 81 resident cooperative CTAs (no + serialized injection row), 4-byte half-plus-generation packets, no sentinel + clearing, half2 projection loads, and a private channel. The original + 40-term FMA chains and cross-warp reduction remain unchanged. Both split and + fused variants compile with 31 registers/16 bytes shared/zero stack or + spills. Compare complete HC with the newly registered up fusion held fixed, + including actual auxiliary sum2 and post-wrap checks. GPU gate pending. diff --git a/tests/distributed/test_custom_all_reduce_dispatch.py b/tests/distributed/test_custom_all_reduce_dispatch.py index 3703931fd3..9c5ee6ca52 100644 --- a/tests/distributed/test_custom_all_reduce_dispatch.py +++ b/tests/distributed/test_custom_all_reduce_dispatch.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock + import pytest import torch +import vllm._custom_ops as ops from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce @@ -33,3 +37,79 @@ def test_should_custom_ar_rejects_unsupported_dtype(dtype: torch.dtype) -> None: communicator = _mock_communicator() assert not communicator.should_custom_ar(torch.empty(16, dtype=dtype)) + + +@pytest.mark.parametrize("sidecar_owner", [False, True]) +@pytest.mark.parametrize("owner_has_op", [False, True]) +@pytest.mark.parametrize( + "op_name,num_inputs", + [("sm70_qwen38_hc_output_allgather", 1), ("sm70_qwen38_hc_up_mix_allgather", 3)], +) +def test_hc_gather_stays_in_communicator_dso( + monkeypatch: pytest.MonkeyPatch, + sidecar_owner: bool, + owner_has_op: bool, + op_name: str, + num_inputs: int, +) -> None: + base = SimpleNamespace(init_custom_ar=Mock()) + sidecar = SimpleNamespace() + if sidecar_owner: + sidecar.init_custom_ar = Mock() + owner, other = (sidecar, base) if sidecar_owner else (base, sidecar) + setattr(other, op_name, Mock()) + if owner_has_op: + setattr(owner, op_name, Mock()) + monkeypatch.setattr(torch.ops, "_C_custom_ar", base) + monkeypatch.setattr(torch.ops, "_C_custom_ar_flashnext", sidecar) + assert getattr(ops, f"supports_{op_name}")() == owner_has_op + tensors = [torch.empty(16) for _ in range(num_inputs + 1)] + if owner_has_op: + getattr(ops, op_name)(123, *tensors) + getattr(owner, op_name).assert_called_once_with(123, *tensors) + else: + # An old sidecar must not borrow the new op from a rebuilt base wheel, + # and a sidecar without init must not receive the base wheel's pointer. + with pytest.raises(AttributeError): + getattr(ops, op_name)(123, *tensors) + getattr(other, op_name).assert_not_called() + + +@pytest.mark.parametrize("fused,hidden", [(True, True), (False, True), (False, False)]) +def test_hc_model_selects_available_owner_route(monkeypatch, fused, hidden) -> None: + import vllm.distributed.parallel_state as parallel + import vllm.models.qwen4_exp.nvidia.sm70_fp16_hc as hc + + comm = SimpleNamespace( + rank=0, + can_sm70_qwen38_hc_shard=Mock(return_value=True), + supports_sm70_qwen38_hc_up_mix_allgather=Mock(return_value=fused), + supports_sm70_qwen38_hc_output_allgather=Mock(return_value=hidden), + sm70_qwen38_hc_down_allgather=Mock(), + sm70_qwen38_hc_up_mix_allgather=Mock(), + sm70_qwen38_hc_output_allgather=Mock(), + sm70_qwen38_hc_gate_mix=Mock(), + ) + tp = SimpleNamespace(device_communicator=SimpleNamespace(ca_comm=comm)) + monkeypatch.setattr(parallel, "get_tp_group", lambda: tp) + monkeypatch.setattr(hc, "_runtime_ok", lambda *args: True) + for name in ( + "_qwen38_hc_down_local_shard_kernel", + "_qwen38_hc_up_hidden_shard_kernel", + "_qwen38_hc_up_local_gate_kernel", + ): + monkeypatch.setattr(hc, name, MagicMock()) + block, injection = hc._qwen38_sm70_fp16_fused_hc( + torch.empty(1, 10240), torch.empty(0), torch.empty(0) + ) + assert block.shape == (1, 2560) and injection.shape == (1, 4) + comm.sm70_qwen38_hc_down_allgather.assert_called_once() + expected = ( + "up_mix_allgather" if fused else "output_allgather" if hidden else "gate_mix" + ) + for name in ("up_mix_allgather", "output_allgather", "gate_mix"): + op = getattr(comm, f"sm70_qwen38_hc_{name}") + if name == expected: + op.assert_called_once() + else: + op.assert_not_called() diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index e6511234d4..0cbcdec81c 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -905,6 +905,28 @@ def test_qsa_lexicographic_topk_is_exact_and_repeatable() -> None: torch.testing.assert_close(output, expected, rtol=0, atol=0) +@pytest.mark.parametrize("live_length", [2048, 2304, 2305]) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@torch.inference_mode() +def test_qsa_lexicographic_topk_decode_boundary_is_exact( + live_length: int, +) -> None: + """The decode fast path and its capacity fallback retain exact ties.""" + + torch.set_default_device("cuda:0") + top_k = 512 + logits = torch.randn((1, 4096), dtype=torch.float32) + logits[0, :live_length:3] = 0.0 + lengths = torch.tensor([live_length], dtype=torch.int32) + output = torch.empty((1, top_k), dtype=torch.int32) + expected = _qsa_lexicographic_topk_reference(logits, [live_length], top_k) + + for _ in range(10): + torch.ops._C.qsa_lexicographic_topk(logits, lengths, output, top_k) + torch.accelerator.synchronize() + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @torch.inference_mode() def test_qsa_lexicographic_topk_supports_prefill_batches() -> None: @@ -935,7 +957,7 @@ def test_qsa_lexicographic_topk_cuda_graph_replay_is_stable() -> None: torch.set_default_device("cuda:0") top_k = 512 - live_length = 4096 + live_length = 2176 logits = torch.zeros((1, 8192), dtype=torch.float32) lengths = torch.tensor([live_length], dtype=torch.int32) output = torch.empty((1, top_k), dtype=torch.int32) diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py index e8a676642a..2abca7927f 100644 --- a/tests/models/qwen4_exp/test_qsa_ops.py +++ b/tests/models/qwen4_exp/test_qsa_ops.py @@ -13,6 +13,7 @@ _qsa_indexer_cublas_shape_supported, _qsa_sparse_launch_profile, _qsa_xqa_page4_shape_supported, + _sm70_qsa_lexicographic_topk_op, _use_sm70_qsa_lexicographic_topk, ) @@ -473,3 +474,21 @@ def test_qsa_lexicographic_topk_is_limited_to_sm70_qsa_shape(monkeypatch): lambda capability: False, ) assert not _use_sm70_qsa_lexicographic_topk(512) + + +def test_qsa_lexicographic_topk_prefers_validation_sidecar(monkeypatch): + sidecar = object() + wheel = object() + monkeypatch.setattr( + qsa_ops.torch, + "ops", + SimpleNamespace( + _C_qsa_sm70=SimpleNamespace(qsa_lexicographic_topk=sidecar), + _C=SimpleNamespace(qsa_lexicographic_topk=wheel), + ), + ) + + assert _sm70_qsa_lexicographic_topk_op() is sidecar + + qsa_ops.torch.ops._C_qsa_sm70 = SimpleNamespace() + assert _sm70_qsa_lexicographic_topk_op() is wheel diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index dc16187b54..db4482c0ae 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -640,6 +640,7 @@ def test_qsa_block_expansion_matches_test_reference() -> None: [ # Kernel-visible pages with --block-size 256 and hybrid-cache alignment. pytest.param(1, 24, 2, 1792, id="tp1_split64"), + pytest.param(1, 6, 1, 1024, id="tp4_decode_split64"), pytest.param(16, 12, 1, 1792, id="tp2_split32"), pytest.param(32, 6, 1, 1024, id="tp4_split8"), pytest.param(257, 6, 1, 1024, id="tp4_split4"), @@ -665,6 +666,7 @@ def test_qsa_sparse_paged_attention_matches_test_reference( q = torch.randn( num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16 ) + output_gate = torch.randn_like(q) kv_cache = torch.randn( num_cache_blocks, page_size, @@ -718,6 +720,14 @@ def test_qsa_sparse_paged_attention_matches_test_reference( assert logical_indices.shape == (num_rows, selection_width) scale = q.shape[-1] ** -0.5 + ungated = qsa_ops.qsa_sparse_paged_attention( + q, + k_cache, + v_cache, + logical_indices, + block_table, + token_to_req, + ) actual = qsa_ops.qsa_sparse_paged_attention( q, k_cache, @@ -725,7 +735,14 @@ def test_qsa_sparse_paged_attention_matches_test_reference( logical_indices, block_table, token_to_req, + output_gate=output_gate, ) + # Match the compiled model path: load the rounded attention and gate in + # FP32, then evaluate sigmoid and multiply before the final BF16 store. + expected_fused = ungated.clone() + qsa_ops._qsa_output_gate(expected_fused, output_gate) + torch.testing.assert_close(actual, expected_fused, rtol=0, atol=0) + expected = _qsa_sparse_paged_attention_reference( q, k_cache, @@ -735,6 +752,7 @@ def test_qsa_sparse_paged_attention_matches_test_reference( token_to_req, scale, ) + expected = expected * torch.sigmoid(output_gate) torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) diff --git a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py index 9a4bfe1e3b..8b33802cc7 100644 --- a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py +++ b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py @@ -2,9 +2,21 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch import vllm.envs as envs +from vllm.models.qwen4_exp.nvidia.ops.hc import hc_gate_mix from vllm.models.qwen4_exp.nvidia.sm70_fp16_gemv import _plan_for +from vllm.models.qwen4_exp.nvidia.sm70_fp16_hc import ( + _qwen38_hc_down_local_shard_kernel, + _qwen38_hc_down_silu_inject_kernel, + _qwen38_hc_up_gate_mix_kernel, + _qwen38_hc_up_gate_mix_row4_kernel, + _qwen38_hc_up_hidden_shard_kernel, + _qwen38_hc_up_local_gate_kernel, +) +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON def test_qwen38_sm70_fp16_gemv_is_opt_in(monkeypatch: pytest.MonkeyPatch) -> None: @@ -69,3 +81,165 @@ def test_qwen38_sm70_fp16_gemv_rejects_other_roles( prefix: str, shape: tuple[int, int] ) -> None: assert _plan_for(prefix, shape) is None + + +@pytest.mark.skipif( + not current_platform.is_device_capability((7, 0)) or not HAS_TRITON, + reason="Qwen3.8 HC row-tile kernel requires CUDA SM70 and Triton", +) +def test_qwen38_sm70_hc_up_row4_is_bitwise() -> None: + lora = torch.empty(1, 320, dtype=torch.float16, device="cuda") + weight = torch.randn(10240, 320, dtype=torch.float16, device="cuda") + branches = torch.empty(1, 10240, dtype=torch.float16, device="cuda") + reference = torch.empty(1, 2560, dtype=torch.float16, device="cuda") + actual = torch.empty_like(reference) + + for seed in range(8): + torch.manual_seed(seed) + lora.normal_() + branches.normal_() + _qwen38_hc_up_gate_mix_kernel[(2560,)]( + lora, + weight, + branches, + reference, + K=320, + HC_DIMENSION=2560, + HC_COUNT=4, + BLOCK_K=512, + num_warps=2, + ) + _qwen38_hc_up_gate_mix_row4_kernel[(640,)]( + lora, + weight, + branches, + actual, + K=320, + HC_DIMENSION=2560, + HC_COUNT=4, + BLOCK_N=4, + BLOCK_K=512, + num_warps=8, + ) + torch.accelerator.synchronize() + assert torch.equal(actual, reference) + + +@pytest.mark.skipif( + not current_platform.is_device_capability((7, 0)) or not HAS_TRITON, + reason="Qwen3.8 HC TP4 shards require CUDA SM70 and Triton", +) +def test_qwen38_sm70_hc_tp4_compute_shards_are_bitwise() -> None: + x = torch.empty(1, 10240, dtype=torch.float16, device="cuda") + down_weight = torch.randn(336, 10240, dtype=torch.float16, device="cuda") + up_weight = torch.randn(10240, 320, dtype=torch.float16, device="cuda") + reference_lora = torch.empty(1, 320, dtype=torch.float16, device="cuda") + reference_injection = torch.empty(1, 4, dtype=torch.float16, device="cuda") + reference_block = torch.empty(1, 2560, dtype=torch.float16, device="cuda") + + for seed in range(4): + torch.manual_seed(seed) + x.normal_() + _qwen38_hc_down_silu_inject_kernel[(324,)]( + x, + down_weight, + reference_lora, + reference_injection, + K=10240, + BLOCK_K=256, + RANK_VALUE=320, + HC_COUNT=4, + num_warps=4, + ) + _qwen38_hc_up_gate_mix_row4_kernel[(640,)]( + reference_lora, + up_weight, + x, + reference_block, + K=320, + HC_DIMENSION=2560, + HC_COUNT=4, + BLOCK_N=4, + BLOCK_K=512, + num_warps=8, + ) + + local_down = [] + local_gates = [] + for rank in range(4): + shard = torch.empty(1, 88, dtype=torch.float16, device="cuda") + _qwen38_hc_down_local_shard_kernel[(88,)]( + x, + down_weight, + shard, + TP_RANK=rank, + num_warps=4, + ) + local_down.append(shard) + gathered_lora = torch.cat([shard[..., :80] for shard in local_down], dim=-1) + gathered_injection = torch.cat( + [shard[..., 80:81] for shard in local_down], dim=-1 + ) + for rank in range(4): + gate = torch.empty(1, 2560, dtype=torch.float16, device="cuda") + _qwen38_hc_up_local_gate_kernel[(320,)]( + gathered_lora, + up_weight, + gate, + TP_RANK=rank, + BLOCK_N=8, + num_warps=8, + ) + local_gates.append(gate) + actual_block = hc_gate_mix(x, torch.cat(local_gates, dim=-1), 4) + torch.accelerator.synchronize() + + assert torch.equal(gathered_lora, reference_lora) + assert torch.equal(gathered_injection, reference_injection) + assert torch.equal(actual_block, reference_block) + + +@pytest.mark.skipif( + not current_platform.is_device_capability((7, 0)) or not HAS_TRITON, + reason="Qwen3.8 HC hidden shards require CUDA SM70 and Triton", +) +def test_qwen38_sm70_hc_up_hidden_shards_are_bitwise() -> None: + generator = torch.Generator(device="cuda").manual_seed(20260905) + weight = torch.randn( + 10240, 320, dtype=torch.float16, device="cuda", generator=generator + ) + lora = torch.empty(1, 320, dtype=torch.float16, device="cuda") + branches = torch.empty(1, 10240, dtype=torch.float16, device="cuda") + expected = torch.empty(1, 2560, dtype=torch.float16, device="cuda") + actual = torch.empty_like(expected) + for case in range(16): + lora.normal_(generator=generator) + branches.normal_(generator=generator) + if case == 0: + lora.zero_() + elif case == 1: + branches.zero_() + elif case == 2: + lora.mul_(0.01) + _qwen38_hc_up_gate_mix_row4_kernel[(640,)]( + lora, + weight, + branches, + expected, + K=320, + HC_DIMENSION=2560, + HC_COUNT=4, + BLOCK_N=4, + BLOCK_K=512, + num_warps=8, + ) + for rank in range(4): + _qwen38_hc_up_hidden_shard_kernel[(320,)]( + lora, + weight, + branches, + actual[..., rank * 640 : (rank + 1) * 640], + TP_RANK=rank, + num_warps=8, + ) + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 4b50a35d42..aeaf622eee 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -17,6 +17,7 @@ ModelOptNvFp4Config, ) from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + _QWEN38_QPN_M1_W13_SPLIT_K, ModelOptNvFp4SM70MoEMethod, _mtp_weighted_reduce, _prepare_compact_slot_groups, @@ -31,6 +32,10 @@ ) +def test_qwen38_qpn_m1_w13_uses_same_precision_split16_plan(): + assert _QWEN38_QPN_M1_W13_SPLIT_K == 16 + + @pytest.mark.parametrize( ("top_k", "compact_tokens", "dense_tokens"), [(8, 10, 11), (10, 8, 9)], @@ -119,6 +124,19 @@ def test_qwen38_fast_prefill_defaults_on_and_can_be_disabled(monkeypatch): monkeypatch.setenv(name, "0") assert not envs.VLLM_SM70_NVFP4_QWEN38_MOE_FUSED_SWIGLU_PREFILL + +def test_qwen38_w2_direct_reduce_defaults_on_and_can_be_disabled(monkeypatch): + name = "VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + try: + assert envs.VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE + monkeypatch.setenv(name, "0") + envs.disable_envs_cache() + assert not envs.VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE + finally: + envs.disable_envs_cache() + name = "VLLM_SM70_NVFP4_QWEN38_MOE_FAST_PREFILL" monkeypatch.delenv(name, raising=False) assert envs.VLLM_SM70_NVFP4_QWEN38_MOE_FAST_PREFILL diff --git a/tests/quantization/test_sm70_online_qpn8.py b/tests/quantization/test_sm70_online_qpn8.py index d72eed9892..e5090c3c19 100644 --- a/tests/quantization/test_sm70_online_qpn8.py +++ b/tests/quantization/test_sm70_online_qpn8.py @@ -77,6 +77,28 @@ def test_nvfp4_mtp5_capability_is_not_inferred_from_m1(monkeypatch): assert online_qpn8.sm70_ops.has_nvfp4_qpn_mtp5_dispatch() +def test_nvfp4_w2_direct_reduce_capability_is_explicit(monkeypatch): + legacy_sidecar = SimpleNamespace(nvfp4_moe_qpn_m1_sm70_out=object()) + monkeypatch.setattr(torch.ops, "_C_qwen38", legacy_sidecar) + monkeypatch.setattr(torch.ops, "_C", SimpleNamespace()) + + assert not online_qpn8.sm70_ops.has_nvfp4_qwen38_w2_direct_reduce() + + legacy_sidecar.nvfp4_qwen38_w2_direct_reduce_out = object() + assert online_qpn8.sm70_ops.has_nvfp4_qwen38_w2_direct_reduce() + + +def test_qwen38_shared_gate_exact_capability_is_explicit(monkeypatch): + sidecar = SimpleNamespace(qwen38_shared_gate_exact_out=object()) + monkeypatch.setattr(torch.ops, "_C_qwen38", sidecar) + monkeypatch.setattr(torch.ops, "_C", SimpleNamespace()) + + assert online_qpn8.sm70_ops.has_qwen38_shared_gate_exact() + + del sidecar.qwen38_shared_gate_exact_out + assert not online_qpn8.sm70_ops.has_qwen38_shared_gate_exact() + + @pytest.mark.parametrize( ("prefix", "k", "n", "expected"), [ diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 85ff6c2d69..ce8669d124 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -231,6 +231,19 @@ def test_sm70_tp4_push_allreduce_mtp5_is_opt_in(monkeypatch): envs.disable_envs_cache() +def test_sm70_tp4_push_allreduce_sum2_m1_is_default_on_with_rollback(monkeypatch): + name = "VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1" + monkeypatch.delenv(name, raising=False) + envs.disable_envs_cache() + try: + assert getattr(envs, name) + monkeypatch.setenv(name, "0") + envs.disable_envs_cache() + assert not getattr(envs, name) + finally: + envs.disable_envs_cache() + + def _bare_dflash2_model() -> DFlash2Qwen3Model: model = DFlash2Qwen3Model.__new__(DFlash2Qwen3Model) torch.nn.Module.__init__(model) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index f2bea0da4b..f8e23d6c7e 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3085,6 +3085,58 @@ def all_reduce_sum2( _custom_ar_op("all_reduce_sum2")(fa, inp_a, inp_b, out) +def sm70_qwen38_hc_down_allgather( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, +) -> None: + _custom_ar_op("sm70_qwen38_hc_down_allgather")(fa, inp, out) + + +def sm70_qwen38_hc_gate_mix( + fa: int, + local_gate: torch.Tensor, + branches: torch.Tensor, + out: torch.Tensor, +) -> None: + _custom_ar_op("sm70_qwen38_hc_gate_mix")(fa, local_gate, branches, out) + + +def _custom_ar_owner_namespace(): + # The opaque communicator belongs to the DSO that initialized it. A new + # optional op must not fall through to another DSO with a different ABI. + sidecar = torch.ops._C_custom_ar_flashnext + return sidecar if hasattr(sidecar, "init_custom_ar") else torch.ops._C_custom_ar + + +def supports_sm70_qwen38_hc_output_allgather() -> bool: + return hasattr(_custom_ar_owner_namespace(), "sm70_qwen38_hc_output_allgather") + + +def sm70_qwen38_hc_output_allgather( + fa: int, + local_block: torch.Tensor, + out: torch.Tensor, +) -> None: + _custom_ar_owner_namespace().sm70_qwen38_hc_output_allgather(fa, local_block, out) + + +def supports_sm70_qwen38_hc_up_mix_allgather() -> bool: + return hasattr(_custom_ar_owner_namespace(), "sm70_qwen38_hc_up_mix_allgather") + + +def sm70_qwen38_hc_up_mix_allgather( + fa: int, + lora: torch.Tensor, + weight: torch.Tensor, + branches: torch.Tensor, + out: torch.Tensor, +) -> None: + _custom_ar_owner_namespace().sm70_qwen38_hc_up_mix_allgather( + fa, lora, weight, branches, out + ) + + def top1_argmax( fa: int, input_pair: torch.Tensor, diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 6fc4440ca1..fe85479fb7 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -154,6 +154,24 @@ def has_nvfp4_qpn_m1_dispatch() -> bool: ) +def has_nvfp4_qwen38_w2_direct_reduce() -> bool: + return hasattr(torch.ops._C_qwen38, "nvfp4_qwen38_w2_direct_reduce_out") or hasattr( + torch.ops._C, "nvfp4_qwen38_w2_direct_reduce_out" + ) + + +def has_nvfp4_qwen38_w13_fused_swiglu() -> bool: + return hasattr(torch.ops._C_qwen38, "nvfp4_qwen38_w13_fused_swiglu_out") or hasattr( + torch.ops._C, "nvfp4_qwen38_w13_fused_swiglu_out" + ) + + +def has_qwen38_shared_gate_exact() -> bool: + return hasattr(torch.ops._C_qwen38, "qwen38_shared_gate_exact_out") or hasattr( + torch.ops._C, "qwen38_shared_gate_exact_out" + ) + + def has_nvfp4_qpn_mtp5_dispatch() -> bool: """Reject extensions that only implement the legacy ten-route kernel.""" return hasattr(torch.ops._C_qwen38, "nvfp4_moe_qpn_mtp5_sm70_out") or hasattr( @@ -1542,6 +1560,115 @@ def _nvfp4_moe_qpn_m1_sm70_out_sidecar_fake( return None +def nvfp4_qwen38_w2_direct_reduce_out( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + topk_weights: torch.Tensor, +) -> None: + _qwen38_qpn8_op("nvfp4_qwen38_w2_direct_reduce_out")( + out, input, weights, scales, expert_ids, topk_weights + ) + + +if hasattr(torch.ops._C, "nvfp4_qwen38_w2_direct_reduce_out"): + + @register_fake("_C::nvfp4_qwen38_w2_direct_reduce_out") + def _nvfp4_qwen38_w2_direct_reduce_out_fake( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> None: + return None + + +if hasattr(torch.ops._C_qwen38, "nvfp4_qwen38_w2_direct_reduce_out"): + + @register_fake("_C_qwen38::nvfp4_qwen38_w2_direct_reduce_out") + def _nvfp4_qwen38_w2_direct_reduce_out_sidecar_fake( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> None: + return None + + +def nvfp4_qwen38_w13_fused_swiglu_out( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, +) -> None: + _qwen38_qpn8_op("nvfp4_qwen38_w13_fused_swiglu_out")( + out, input, weights, scales, expert_ids + ) + + +if hasattr(torch.ops._C, "nvfp4_qwen38_w13_fused_swiglu_out"): + + @register_fake("_C::nvfp4_qwen38_w13_fused_swiglu_out") + def _nvfp4_qwen38_w13_fused_swiglu_out_fake( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + ) -> None: + return None + + +if hasattr(torch.ops._C_qwen38, "nvfp4_qwen38_w13_fused_swiglu_out"): + + @register_fake("_C_qwen38::nvfp4_qwen38_w13_fused_swiglu_out") + def _nvfp4_qwen38_w13_fused_swiglu_out_sidecar_fake( + out: torch.Tensor, + input: torch.Tensor, + weights: torch.Tensor, + scales: torch.Tensor, + expert_ids: torch.Tensor, + ) -> None: + return None + + +def qwen38_shared_gate_exact_out( + out: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, +) -> None: + _qwen38_qpn8_op("qwen38_shared_gate_exact_out")(out, input, weight) + + +if hasattr(torch.ops._C, "qwen38_shared_gate_exact_out"): + + @register_fake("_C::qwen38_shared_gate_exact_out") + def _qwen38_shared_gate_exact_out_fake( + out: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + ) -> None: + return None + + +if hasattr(torch.ops._C_qwen38, "qwen38_shared_gate_exact_out"): + + @register_fake("_C_qwen38::qwen38_shared_gate_exact_out") + def _qwen38_shared_gate_exact_out_sidecar_fake( + out: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + ) -> None: + return None + + def nvfp4_moe_qpn_mtp5_sm70_out( out: torch.Tensor, input: torch.Tensor, diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 3bf6f6cf31..2c8d6ffdce 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -340,10 +340,15 @@ def __init__( mtp5_status = ( "enabled" if envs.VLLM_SM70_TP4_PUSH_ALLREDUCE_MTP5 else "disabled" ) + sum2_m1_status = ( + "enabled" if envs.VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1 else "disabled" + ) logger.info( "SM70 TP4 SGLang-style push all-reduce enabled for the " "FP16 80-KiB verifier, 8-KiB decode, and 5-KiB Qwen4Exp " - "payloads; opt-in 25-KiB Qwen4Exp MTP4 payload is %s.", + "payloads; 5-KiB Qwen4Exp sum2 is %s and opt-in 25-KiB " + "Qwen4Exp MTP4 is %s.", + sum2_m1_status, mtp5_status, ) @@ -435,6 +440,51 @@ def all_reduce_sum2( ops.all_reduce_sum2(self._ptr, inp_a, inp_b, out) return out + def can_sm70_qwen38_hc_shard(self, branches: torch.Tensor) -> 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 branches.is_cuda + and branches.dtype == torch.float16 + and branches.shape == (1, 10240) + and branches.is_contiguous() + ) + + def sm70_qwen38_hc_down_allgather( + self, local_down: torch.Tensor, gathered_down: torch.Tensor + ) -> None: + ops.sm70_qwen38_hc_down_allgather(self._ptr, local_down, gathered_down) + + def sm70_qwen38_hc_gate_mix( + self, + local_gate: torch.Tensor, + branches: torch.Tensor, + output: torch.Tensor, + ) -> None: + ops.sm70_qwen38_hc_gate_mix(self._ptr, local_gate, branches, output) + + def supports_sm70_qwen38_hc_output_allgather(self) -> bool: + return ops.supports_sm70_qwen38_hc_output_allgather() + + def sm70_qwen38_hc_output_allgather( + self, local_block: torch.Tensor, output: torch.Tensor + ) -> None: + ops.sm70_qwen38_hc_output_allgather(self._ptr, local_block, output) + + def supports_sm70_qwen38_hc_up_mix_allgather(self) -> bool: + return ops.supports_sm70_qwen38_hc_up_mix_allgather() + + def sm70_qwen38_hc_up_mix_allgather( + self, + lora: torch.Tensor, + weight: torch.Tensor, + branches: torch.Tensor, + output: torch.Tensor, + ) -> None: + ops.sm70_qwen38_hc_up_mix_allgather(self._ptr, lora, weight, branches, output) + def sm70_tp2_all_reduce_gemma_rms_norm( self, inp: torch.Tensor, diff --git a/vllm/envs.py b/vllm/envs.py index d50e8dd634..f277420af5 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -186,6 +186,7 @@ VLLM_SM70_NVFP4_TUNE_SMALL_SHAPES: bool = True VLLM_SM70_NVFP4_QWEN38_TP4_M1_FAST_SELECTOR: bool = True VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE: bool = True + VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE: bool = True VLLM_SM70_NVFP4_QWEN38_MOE_INDEXED_PREFILL: bool = True VLLM_SM70_NVFP4_QWEN38_MOE_FUSED_SWIGLU_PREFILL: bool = True VLLM_SM70_NVFP4_QWEN38_MOE_FAST_PREFILL: bool = True @@ -232,6 +233,7 @@ VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC: bool = False VLLM_SM70_TP4_PUSH_ALLREDUCE: bool = True VLLM_SM70_TP4_PUSH_ALLREDUCE_MTP5: bool = False + VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1: bool = True VLLM_SM70_CUSTOM_AR_LIBRARY: str | None = None VLLM_SM70_TOP1_CUSTOM_AR: bool = False VLLM_SM70_GREEDY_TOKEN_FASTPATH: bool = True @@ -1903,6 +1905,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_NVFP4_QWEN38_MOE_QPN_MTP5_DECODE": lambda: bool( int(os.getenv("VLLM_SM70_NVFP4_QWEN38_MOE_QPN_MTP5_DECODE", "0")) ), + # Exact single-token Qwen3.8 W2 epilogue. Ten expert warps retain the + # established FP16 route rounding and reduce in top-k order with FP32 FMA. + "VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE": lambda: bool( + int(os.getenv("VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE", "1")) + ), "VLLM_SM70_NVFP4_QPN_M1_LIBRARY": lambda: os.getenv( "VLLM_SM70_NVFP4_QPN_M1_LIBRARY" ), @@ -2114,6 +2121,14 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_TP4_PUSH_ALLREDUCE_MTP5": lambda: bool( int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE_MTP5", "0")) ), + # Exact Qwen3.8 single-token MoE payload: FP16 [1, 2560]. Reuse the + # already-registered SM70 TP4 push buffers for all_reduce_sum2 while + # retaining the existing FP16 local sum and rank-ordered FP32 reduction. + # The TP4 CUDA Graph gate is bitwise across all ranks and cuts 48 + # collectives from 0.459 ms to 0.136 ms; explicit 0 is the rollback. + "VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1": lambda: bool( + int(os.getenv("VLLM_SM70_TP4_PUSH_ALLREDUCE_SUM2_M1", "1")) + ), # Optional task-built custom-AR fragment. Operators present in the sidecar # override the production namespace; every other operator falls back. "VLLM_SM70_CUSTOM_AR_LIBRARY": lambda: os.getenv("VLLM_SM70_CUSTOM_AR_LIBRARY"), diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 10915381e9..13592cf777 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -49,7 +49,10 @@ _SUPPORTED_TP_SIZES: Final = (1, 2, 4) _GRAPH_SAFE_MAX_TOKENS: Final = 18 _COMPACT_GROUPED_MAX_SLOTS: Final = 80 -_QWEN38_QPN_M1_W13_SPLIT_K: Final = 8 +# V100 real-shape M=1 tuning favors 16 warps. This retains checkpoint NVFP4, +# FP32 MMA accumulation, and the FP16 W13 output boundary; only the order in +# which the FP32 K partitions are joined changes from the former split-8 plan. +_QWEN38_QPN_M1_W13_SPLIT_K: Final = 16 _QWEN38_QPN_M1_W2_SPLIT_K: Final = 1 _QWEN38_INDEXED_PREFILL_MIN_TOKENS: Final = 128 _QWEN38_QPN_MTP5_W13_SPLIT_K: Final = 4 @@ -431,6 +434,25 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: and not sm70_ops.has_nvfp4_qpn_m1_dispatch() ): missing.append("nvfp4_moe_qpn_m1_sm70_out") + w2_direct_reduce_requested = bool( + envs.VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE + ) + w2_direct_reduce_available = sm70_ops.has_nvfp4_qwen38_w2_direct_reduce() + w2_direct_reduce_explicit = ( + "VLLM_SM70_NVFP4_QWEN38_MOE_W2_DIRECT_REDUCE" in os.environ + ) + if ( + w2_direct_reduce_requested + and not w2_direct_reduce_available + and w2_direct_reduce_explicit + ): + missing.append("nvfp4_qwen38_w2_direct_reduce_out") + elif w2_direct_reduce_requested and not w2_direct_reduce_available: + logger.warning_once( + "The default SM70 Qwen3.8 W2 direct-reduce op is absent from " + "the loaded extension; retaining separate W2 and weighted " + "reduce kernels. Explicit opt-in fails closed." + ) if ( envs.VLLM_SM70_NVFP4_QWEN38_MOE_QPN_MTP5_DECODE and not sm70_ops.has_nvfp4_qpn_mtp5_dispatch() @@ -528,6 +550,20 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: "activation route. Explicit opt-in fails closed." ) fused_swiglu_prefill = bool(fused_swiglu_requested and fused_swiglu_available) + fused_swiglu_decode = bool( + envs.VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE + and fused_swiglu_prefill + and sm70_ops.has_nvfp4_qwen38_w13_fused_swiglu() + ) + if ( + envs.VLLM_SM70_NVFP4_QWEN38_MOE_QPN_M1_DECODE + and fused_swiglu_prefill + and not fused_swiglu_decode + ): + logger.warning_once( + "The SM70 Qwen3.8 fused W13/SwiGLU decode op is absent; " + "retaining separate exact W13 and activation kernels." + ) fast_prefill = bool( fused_swiglu_prefill and envs.VLLM_SM70_NVFP4_QWEN38_MOE_FAST_PREFILL ) @@ -645,7 +681,11 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: indexed_prefill_requested and indexed_prefill_available ) layer.sm70_nvfp4_qwen38_fused_swiglu_prefill = fused_swiglu_prefill + layer.sm70_nvfp4_qwen38_fused_swiglu_decode = fused_swiglu_decode layer.sm70_nvfp4_qwen38_fast_prefill = fast_prefill + layer.sm70_nvfp4_qwen38_w2_direct_reduce = bool( + w2_direct_reduce_requested and w2_direct_reduce_available + ) layer.sm70_nvfp4_graph_safe_max_tokens = _GRAPH_SAFE_MAX_TOKENS layer.sm70_nvfp4_compact_grouped_max_slots = _COMPACT_GROUPED_MAX_SLOTS self._allocate_graph_safe_decode_buffers(layer) @@ -674,6 +714,11 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: "SM70 Qwen3.8 indexed-A fused-SwiGLU prefill candidate " "enabled (interleaved W13, exact FP16 epilogue arithmetic)." ) + if fused_swiglu_decode: + logger.info_once( + "SM70 Qwen3.8 fused W13/SwiGLU decode route enabled " + "(split16, exact FP16 rounding and activation arithmetic)." + ) if fast_prefill: logger.info_once( "SM70 Qwen3.8 NVFP4 fast grouped prefill enabled " @@ -944,21 +989,51 @@ def apply( if direct_qpn_m1 else sm70_ops.nvfp4_moe_qpn_mtp5_sm70_out ) - direct_op( - buffers["gate_up"], - x, - layer.w13_tm_weight, - layer.w13_tm_scales, - route_ids, - True, - w13_split_k, - ) - self._apply_swiglu( - layer, - buffers["intermediate"], - buffers["gate_up"], - interleaved=interleaved_w13, + fused_w13_decode = bool( + direct_qpn_m1 + and interleaved_w13 + and getattr( + layer, + "sm70_nvfp4_qwen38_fused_swiglu_decode", + False, + ) ) + if fused_w13_decode: + sm70_ops.nvfp4_qwen38_w13_fused_swiglu_out( + buffers["intermediate"], + x, + layer.w13_tm_weight, + layer.w13_tm_scales, + route_ids, + ) + else: + direct_op( + buffers["gate_up"], + x, + layer.w13_tm_weight, + layer.w13_tm_scales, + route_ids, + True, + w13_split_k, + ) + self._apply_swiglu( + layer, + buffers["intermediate"], + buffers["gate_up"], + interleaved=interleaved_w13, + ) + if direct_qpn_m1 and bool( + getattr(layer, "sm70_nvfp4_qwen38_w2_direct_reduce", False) + ): + sm70_ops.nvfp4_qwen38_w2_direct_reduce_out( + output, + buffers["intermediate"], + layer.w2_tm_weight, + layer.w2_tm_scales, + route_ids, + topk_weights, + ) + return output direct_op( buffers["sorted_output"], buffers["intermediate"], diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index 31ab7968a8..0600787276 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -143,7 +143,7 @@ def __init__( enforce_enable=_sm70_force_shared_expert_silu_custom_op(prefix) ) self.expert_gate = expert_gate - self._sm70_fused_shared_expert_gate = ( + self._sm70_exact_shared_expert_gate = ( envs.VLLM_SM70_QWEN3NEXT_SHARED_GATE_FUSION and expert_gate is not None and prefix.endswith(".mlp.shared_expert") @@ -151,7 +151,6 @@ def __init__( and intermediate_size == 160 and not reduce_results and _sm70_force_shared_expert_silu_custom_op(prefix) - and hasattr(torch.ops._C, "sm70_f16_gate_mul_out") ) def forward(self, x): @@ -166,26 +165,30 @@ def forward(self, x): out, _ = self.down_proj(out) out = _sm70_dump_qwen_mlp_tensor("mlp_down_out", self.layer_idx, out) + used_exact_gate = False if ( - self._sm70_fused_shared_expert_gate + self._sm70_exact_shared_expert_gate and x.shape[0] == 1 and x.dtype == torch.float16 and out.dtype == torch.float16 ): from vllm import _sm70_ops as sm70_ops - assert self.expert_gate is not None - gate_weight = self.expert_gate.weight - if gate_weight.dtype != torch.float16 or not gate_weight.is_contiguous(): - raise RuntimeError( - "SM70 Qwen3Next fused shared-expert gate requires a " - "contiguous FP16 gate weight." - ) - sm70_ops.sm70_f16_gate_mul_out(out, x, gate_weight) - out = _sm70_dump_qwen_mlp_tensor( - "mlp_after_expert_gate", self.layer_idx, out - ) - elif self.expert_gate is not None: + if sm70_ops.has_qwen38_shared_gate_exact(): + assert self.expert_gate is not None + gate_weight = self.expert_gate.weight + if ( + gate_weight.dtype != torch.float16 + or not gate_weight.is_contiguous() + ): + raise RuntimeError( + "SM70 Qwen3.8 exact shared-expert gate requires a " + "contiguous FP16 gate weight." + ) + logger.info_once("SM70 Qwen3.8 exact shared-expert gate path enabled.") + sm70_ops.qwen38_shared_gate_exact_out(out, x, gate_weight) + used_exact_gate = True + if self.expert_gate is not None and not used_exact_gate: expert_gate = self.expert_gate(x)[0] expert_gate = _sm70_dump_qwen_mlp_tensor( "mlp_expert_gate", self.layer_idx, expert_gate @@ -199,6 +202,11 @@ def forward(self, x): "mlp_after_expert_gate", self.layer_idx, out ) + if used_exact_gate: + out = _sm70_dump_qwen_mlp_tensor( + "mlp_after_expert_gate", self.layer_idx, out + ) + return out diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index b8b370ba18..8c4e36c1dd 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -21,6 +21,23 @@ _LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024 _TOPK_WORKSPACE_BYTES = 1024 * 1024 +_SM70_QSA_TOPK_LIBRARY = os.getenv("VLLM_SM70_QSA_TOPK_LIBRARY") +if _SM70_QSA_TOPK_LIBRARY is not None: + torch.ops.load_library(_SM70_QSA_TOPK_LIBRARY) + +if hasattr(torch.ops._C_qsa_sm70, "qsa_lexicographic_topk"): + + @torch.library.register_fake("_C_qsa_sm70::qsa_lexicographic_topk") + def _qsa_lexicographic_topk_sidecar_fake( + logits: torch.Tensor, + lengths: torch.Tensor, + output: torch.Tensor, + topk: int, + ) -> None: + del logits, lengths, output, topk + return None + + _SM70_INDEXER_CUBLAS = os.getenv("VLLM_SM70_QSA_INDEXER_CUBLAS", "1") == "1" _SM70_INDEXER_SCORE_TILE_BYTES = ( int(os.getenv("VLLM_SM70_QSA_INDEXER_SCORE_TILE_MB", "64")) * 1024 * 1024 @@ -523,6 +540,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_q_row, stride_q_head, stride_k_block, @@ -535,6 +553,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( stride_table_req, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, num_cache_blocks, num_requests, @@ -660,6 +680,21 @@ def _qsa_sparse_paged_gqa_splitk_kernel( # V dequantization is linear, so apply its scalar after the # normalized FP32 accumulation instead of to every loaded value. normalized_output *= v_scale + if output_gate_ptr is not None: + # Preserve the compiled path's rounded attention output before + # evaluating the sigmoid gate and final product in FP32. + normalized_output = normalized_output.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + (first_head + head_offsets[:, None]) * stride_output_gate_head + + dim_offsets[None, :], + mask=output_mask, + other=0.0, + ).to(tl.float32) + normalized_output = normalized_output.to(tl.float32) * tl.sigmoid( + output_gate + ) tl.store( output_ptr + row * stride_output_row @@ -701,8 +736,11 @@ def _qsa_merge_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, v_scale, HEAD_DIM: tl.constexpr, @@ -740,12 +778,62 @@ def _qsa_merge_splitk_kernel( # Apply the V scale once after combining all independently normalized # splits. Scaling partials earlier would repeat this work per split. merged *= v_scale + if output_gate_ptr is not None: + merged = merged.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + head * stride_output_gate_head + + dim_offsets + ).to(tl.float32) + merged = merged.to(tl.float32) * tl.sigmoid(output_gate) tl.store( output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, merged, ) +@triton.jit +def _qsa_output_gate_kernel( + output_ptr, + output_gate_ptr, + stride_output_row, + stride_output_head, + stride_output_gate_row, + stride_output_gate_head, + HEAD_DIM: tl.constexpr, +) -> None: + row = tl.program_id(0) + head = tl.program_id(1) + dim_offsets = tl.arange(0, HEAD_DIM) + output = tl.load( + output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets + ).to(tl.float32) + gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + head * stride_output_gate_head + + dim_offsets + ).to(tl.float32) + tl.store( + output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, + output * tl.sigmoid(gate), + ) + + +def _qsa_output_gate(output: torch.Tensor, output_gate: torch.Tensor) -> None: + _qsa_output_gate_kernel[(output.shape[0], output.shape[1])]( + output, + output_gate, + output.stride(0), + output.stride(1), + output_gate.stride(0), + output_gate.stride(1), + HEAD_DIM=output.shape[2], + num_warps=4, + ) + + @triton.jit def _store_qsa_rows_kernel( cache_ptr, @@ -1072,6 +1160,15 @@ def _use_sm70_qsa_lexicographic_topk(topk: int) -> bool: return topk == 512 and current_platform.is_device_capability(70) +def _sm70_qsa_lexicographic_topk_op(): + """Prefer an opt-in source-validation fragment over the wheel op.""" + + sidecar = torch.ops._C_qsa_sm70 + if hasattr(sidecar, "qsa_lexicographic_topk"): + return sidecar.qsa_lexicographic_topk + return torch.ops._C.qsa_lexicographic_topk + + def _qsa_visible_blocks( token_to_req: torch.Tensor, query_positions: torch.Tensor, @@ -1332,7 +1429,7 @@ def qsa_select_paged_tokens( "Using exact SM70 QSA lexicographic top-k " "(score descending, block index ascending)." ) - torch.ops._C.qsa_lexicographic_topk( + _sm70_qsa_lexicographic_topk_op()( logits, visible_blocks, blocks, @@ -1917,6 +2014,7 @@ def qsa_sparse_paged_attention( block_table: torch.Tensor, token_to_req: torch.Tensor, out: torch.Tensor | None = None, + output_gate: torch.Tensor | None = None, query_positions: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, kv_cache_dtype: str = "auto", @@ -1970,6 +2068,12 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse output must match its query") assert out.dtype == q.dtype and out.device == q.device assert out.stride(2) == 1 + output_gate_view = output_gate.view_as(q) if output_gate is not None else None + if output_gate_view is not None: + if output_gate_view.dtype != q.dtype or output_gate_view.device != q.device: + raise ValueError("QSA output gate must match the query dtype and device") + if output_gate_view.stride(2) != 1: + raise ValueError("QSA output gate must be contiguous in head dimension") if not q.shape[0]: return out @@ -1999,6 +2103,8 @@ def qsa_sparse_paged_attention( v_scale, ) if xqa_output is not None: + if output_gate_view is not None: + _qsa_output_gate(xqa_output, output_gate_view) return xqa_output group_size = q.shape[1] // k_cache.shape[2] @@ -2052,6 +2158,7 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, q.stride(0), q.stride(1), k_cache.stride(0), @@ -2064,6 +2171,8 @@ def qsa_sparse_paged_attention( block_table.stride(0), out.stride(0), out.stride(1), + output_gate_view.stride(0) if output_gate_view is not None else 0, + output_gate_view.stride(1) if output_gate_view is not None else 0, q.shape[0], k_cache.shape[0], block_table.shape[0], @@ -2090,8 +2199,11 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, out.stride(0), out.stride(1), + output_gate_view.stride(0) if output_gate_view is not None else 0, + output_gate_view.stride(1) if output_gate_view is not None else 0, q.shape[0], v_scale, HEAD_DIM=q.shape[2], diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index dc6443cfc6..5f12fd53ef 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -188,6 +188,140 @@ def _dequantize_ple_fp8_bytes_kernel( tl.store(output_ptr + offsets, values, mask=mask) +@triton.jit +def _qwen38_ple_m1_ngram_ids_kernel( + input_ids_ptr, + ngram_context_ptr, + multipliers_ptr, + sizes_ptr, + offsets_ptr, + output_ptr, + EOS_TOKEN_ID: tl.constexpr, +): + """Compute the exact Qwen3.8 M=1 ngram-2/3 IDs in one launch.""" + + head = tl.arange(0, 16) + current = tl.load(input_ids_ptr).to(tl.int64) + older = tl.load(ngram_context_ptr).to(tl.int64) + previous = tl.load(ngram_context_ptr + 1).to(tl.int64) + + # ``compute_ngram_ids`` resets history at EOS. The immediately previous + # token remains the ngram-2 source (and is itself EOS), while ngram-3 must + # not reach across that boundary. + older = tl.where(previous == EOS_TOKEN_ID, EOS_TOKEN_ID, older) + multiplier0 = tl.load(multipliers_ptr).to(tl.int64) + multiplier1 = tl.load(multipliers_ptr + 1).to(tl.int64) + multiplier2 = tl.load(multipliers_ptr + 2).to(tl.int64) + mixed2 = (current * multiplier0) ^ (previous * multiplier1) + mixed3 = mixed2 ^ (older * multiplier2) + mixed = tl.where(head < 8, mixed2, mixed3) + + size = tl.load(sizes_ptr + head).to(tl.int64) + offset = tl.load(offsets_ptr + head).to(tl.int64) + remainder = mixed % size + # PTX signed remainder follows the dividend, whereas torch.remainder is + # always non-negative for these positive vocabulary sizes. + remainder = tl.where(remainder < 0, remainder + size, remainder) + tl.store(output_ptr + head, remainder + offset) + + +@triton.jit +def _qwen38_ple_m1_short_conv_kernel( + x_ptr, + state_ptr, + weight_ptr, + output_ptr, + state_index_ptr, + has_initial_ptr, + STATE_STRIDE_0: tl.constexpr, + STATE_STRIDE_1: tl.constexpr, + STATE_STRIDE_2: tl.constexpr, + HAS_INITIAL: tl.constexpr, + NULL_STATE_ID: tl.constexpr, + HIDDEN_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + """Fuse exact Qwen3.8 M=1 dilated conv and state-cache update.""" + + hidden = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + hidden_mask = hidden < HIDDEN_SIZE + state_index = tl.load(state_index_ptr).to(tl.int64) + valid = state_index != NULL_STATE_ID + safe_state_index = tl.where(valid, state_index, 0) + use_initial = valid + if HAS_INITIAL: + use_initial &= tl.load(has_initial_ptr).to(tl.int1) + state_base = safe_state_index * STATE_STRIDE_0 + hidden * STATE_STRIDE_1 + state_mask = hidden_mask & use_initial + + s0 = tl.load(state_ptr + state_base, mask=state_mask, other=0.0).to(tl.float32) + s1 = tl.load( + state_ptr + state_base + STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + s2 = tl.load( + state_ptr + state_base + 2 * STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + s3 = tl.load( + state_ptr + state_base + 3 * STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + s4 = tl.load( + state_ptr + state_base + 4 * STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + s5 = tl.load( + state_ptr + state_base + 5 * STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + s6 = tl.load( + state_ptr + state_base + 6 * STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + s7 = tl.load( + state_ptr + state_base + 7 * STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + s8 = tl.load( + state_ptr + state_base + 8 * STATE_STRIDE_2, + mask=state_mask, + other=0.0, + ).to(tl.float32) + x = tl.load(x_ptr + hidden, mask=hidden_mask, other=0.0).to(tl.float32) + w0 = tl.load(weight_ptr + hidden * 4).to(tl.float32) + w1 = tl.load(weight_ptr + hidden * 4 + 1).to(tl.float32) + w2 = tl.load(weight_ptr + hidden * 4 + 2).to(tl.float32) + w3 = tl.load(weight_ptr + hidden * 4 + 3).to(tl.float32) + conv = s0 * w0 + conv += s3 * w1 + conv += s6 * w2 + conv += x * w3 + # Preserve the depthwise-conv FP16 output boundary. The caller deliberately + # retains native F.silu because its SM70 rounding differs slightly from + # Triton's sigmoid approximation. + conv = conv.to(tl.float16) + tl.store(output_ptr + hidden, tl.where(valid, conv, 0.0), mask=hidden_mask) + + update_mask = hidden_mask & valid + tl.store(state_ptr + state_base, s1, mask=update_mask) + tl.store(state_ptr + state_base + STATE_STRIDE_2, s2, mask=update_mask) + tl.store(state_ptr + state_base + 2 * STATE_STRIDE_2, s3, mask=update_mask) + tl.store(state_ptr + state_base + 3 * STATE_STRIDE_2, s4, mask=update_mask) + tl.store(state_ptr + state_base + 4 * STATE_STRIDE_2, s5, mask=update_mask) + tl.store(state_ptr + state_base + 5 * STATE_STRIDE_2, s6, mask=update_mask) + tl.store(state_ptr + state_base + 6 * STATE_STRIDE_2, s7, mask=update_mask) + tl.store(state_ptr + state_base + 7 * STATE_STRIDE_2, s8, mask=update_mask) + tl.store(state_ptr + state_base + 8 * STATE_STRIDE_2, x, mask=update_mask) + + def _splitmix64(value: int) -> int: value = (value + _SPLITMIX_GAMMA) & _MASK64 value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64 @@ -697,8 +831,7 @@ def compute_ngram_ids( ngram_context: torch.Tensor, ) -> torch.Tensor: """Compute PLE indices for the current, unpadded request layout.""" - input_ids = input_ids.reshape(-1).long() - query_start_loc = query_start_loc.long() + input_ids = input_ids.reshape(-1) num_reqs = query_start_loc.numel() - 1 num_tokens = input_ids.shape[0] @@ -715,6 +848,47 @@ def compute_ngram_ids( if num_reqs <= 0: raise ValueError("PLE requires at least one request") + if ( + not is_offload_process() + and input_ids.is_cuda + and current_platform.is_device_capability((7, 0)) + and num_tokens == 1 + and num_reqs == 1 + and input_ids.dtype in (torch.int32, torch.int64) + and self.ngram_size == 3 + and self.heads_per_ngram == 8 + and self.ngram_heads == 16 + and ngram_context.ndim == 2 + and ngram_context.shape[0] >= 1 + and ngram_context.shape[1] == 2 + and ngram_context.is_cuda + and ngram_context.is_contiguous() + and self.layer_multipliers.is_cuda + and self.ngram_heads_vocab_sizes.is_cuda + and self.ngram_heads_offsets.is_cuda + and input_ids.device + == ngram_context.device + == self.layer_multipliers.device + == self.ngram_heads_vocab_sizes.device + == self.ngram_heads_offsets.device + ): + output = torch.empty((1, 16), dtype=torch.long, device=input_ids.device) + _qwen38_ple_m1_ngram_ids_kernel[(1,)]( + input_ids, + ngram_context, + self.layer_multipliers, + self.ngram_heads_vocab_sizes, + self.ngram_heads_offsets, + output, + EOS_TOKEN_ID=self.eos_token_id, + num_warps=1, + ) + logger.info_once("SM70 Qwen3.8 fused M=1 PLE ngram-ID path enabled.") + return output + + input_ids = input_ids.long() + query_start_loc = query_start_loc.long() + if is_offload_process(): max_seq_len = max( 1, @@ -1222,6 +1396,67 @@ def _short_conv_dilated_decode_batched( state_indices_tensor_d: torch.Tensor, has_initial_states_d: torch.Tensor | None, ) -> torch.Tensor: + has_initial_ok = has_initial_states_d is None or ( + has_initial_states_d.numel() >= 1 + and has_initial_states_d.is_cuda + and has_initial_states_d.is_contiguous() + ) + if ( + current_platform.is_device_capability((7, 0)) + and x_d.shape == (1, 10240) + and x_d.dtype == torch.float16 + and x_d.is_cuda + and x_d.is_contiguous() + and conv_state.ndim == 3 + and conv_state.shape[1] == 10240 + and conv_state.shape[2] == 9 + and conv_state.dtype == torch.float16 + and conv_state.is_cuda + and conv_weights.shape == (10240, 4) + and conv_weights.dtype == torch.float16 + and conv_weights.is_cuda + and conv_weights.is_contiguous() + and state_indices_tensor_d.numel() == 1 + and state_indices_tensor_d.dtype in (torch.int32, torch.int64) + and state_indices_tensor_d.is_cuda + and state_indices_tensor_d.is_contiguous() + and has_initial_ok + and x_d.device + == conv_state.device + == conv_weights.device + == state_indices_tensor_d.device + and ( + has_initial_states_d is None + or has_initial_states_d.device == x_d.device + ) + ): + conv_output = torch.empty_like(x_d) + has_initial_ptr = ( + state_indices_tensor_d + if has_initial_states_d is None + else has_initial_states_d + ) + _qwen38_ple_m1_short_conv_kernel[(triton.cdiv(10240, 256),)]( + x_d, + conv_state, + conv_weights, + conv_output, + state_indices_tensor_d, + has_initial_ptr, + STATE_STRIDE_0=conv_state.stride(0), + STATE_STRIDE_1=conv_state.stride(1), + STATE_STRIDE_2=conv_state.stride(2), + HAS_INITIAL=has_initial_states_d is not None, + NULL_STATE_ID=NULL_BLOCK_ID, + HIDDEN_SIZE=10240, + BLOCK=256, + num_warps=4, + ) + logger.info_once( + "SM70 Qwen3.8 fused M=1 PLE short-conv state path enabled." + ) + return F.silu(conv_output) + state_indices = state_indices_tensor_d.to( device=conv_state.device, dtype=torch.int64 ) diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index c0d4147e1f..6e5a0c26af 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -145,6 +145,7 @@ def forward_qsa( attn_metadata: FlashAttentionMetadata, output: torch.Tensor, token_to_req: torch.Tensor, + output_gate: torch.Tensor | None = None, query_positions: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, output_scale: torch.Tensor | None = None, @@ -188,6 +189,8 @@ def forward_qsa( qsa_metadata["query_positions"] = query_positions[:num_tokens] if sequence_lengths is not None: qsa_metadata["sequence_lengths"] = sequence_lengths + if output_gate is not None: + qsa_metadata["output_gate"] = output_gate[:num_tokens] qsa_sparse_paged_attention( query[:num_tokens], key_cache, @@ -514,6 +517,7 @@ def _run_qsa( key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, + output_gate: torch.Tensor | None = None, ) -> None: if not self._qsa_kv_scales_finalized: raise RuntimeError( @@ -577,6 +581,7 @@ def _run_qsa( main_metadata, output, token_to_req=side_metadata.token_to_req, + output_gate=output_gate, query_positions=side_metadata.logical_positions, sequence_lengths=side_metadata.seq_lens, ) @@ -609,6 +614,7 @@ def forward( key, value, attn_output, + gate, encoded_layer_name, ) else: @@ -619,11 +625,10 @@ def forward( key, value, attn_output, + gate, encoded_layer_name, ) flat_output = attn_output.view(num_tokens, -1) - if gate is not None: - flat_output = flat_output * torch.sigmoid(gate) projected_output, _ = self.o_proj(flat_output) if output is not None: output.copy_(projected_output) @@ -637,6 +642,7 @@ def qwen4_exp_qsa_with_output( key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, + output_gate: torch.Tensor | None, layer_name: LayerNameType, ) -> None: """Run the complete QSA state/update/attend transaction.""" @@ -652,6 +658,7 @@ def qwen4_exp_qsa_with_output( key, value, output, + output_gate, ) @@ -662,9 +669,10 @@ def qwen4_exp_qsa_with_output_fake( key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, + output_gate: torch.Tensor | None, layer_name: LayerNameType, ) -> None: - del hidden_states, positions, query, key, value, output, layer_name + del hidden_states, positions, query, key, value, output, output_gate, layer_name direct_register_custom_op( diff --git a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py index 4574bf3d6f..b271461663 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py +++ b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py @@ -63,6 +63,111 @@ def _qwen38_hc_down_silu_inject_kernel( tl.store(injection_ptr + row - RANK_VALUE, value, mask=~is_lora) +@triton.jit +def _qwen38_hc_down_local_shard_kernel( + x_ptr, + weight_ptr, + output_ptr, + TP_RANK: tl.constexpr, +): + """Compute this TP rank's 80 low-rank rows and one injection row.""" + row = tl.program_id(0) + active = row < 81 + checkpoint_row = tl.where(row < 80, TP_RANK * 80 + row, 320 + TP_RANK) + offsets = tl.arange(0, 256) + acc = tl.zeros((256,), dtype=tl.float32) + for block_start in tl.static_range(0, 10240, 256): + indices = block_start + offsets + x = tl.load( + x_ptr + indices, + mask=active, + other=0.0, + eviction_policy="evict_last", + ) + weight = tl.load( + weight_ptr + checkpoint_row * 10240 + indices, + mask=active, + other=0.0, + eviction_policy="evict_first", + ) + acc += x.to(tl.float32) * weight.to(tl.float32) + + # Match the replicated projection's FP16 materialization before SiLU. + value = tl.sum(acc, axis=0).to(tl.float16).to(tl.float32) + scaled = value / 4 + value = tl.where(row < 80, scaled * tl.sigmoid(scaled), value) + tl.store(output_ptr + row, value, mask=active) + # Keep the 88-element communication packet aligned to 16 bytes. Padding + # is canonical zero and is discarded after the rank-ordered gather. + tl.store(output_ptr + row, 0.0, mask=~active) + + +@triton.jit +def _qwen38_hc_up_local_gate_kernel( + lora_ptr, + weight_ptr, + gate_ptr, + TP_RANK: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Compute the 2560 gate rows owned by this TP rank.""" + hidden = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + offsets = tl.arange(0, 512) + hidden_mask = hidden < 2560 + k_mask = offsets < 320 + lora = tl.load( + lora_ptr + offsets, + mask=k_mask, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + checkpoint_row = TP_RANK * 2560 + hidden + weight = tl.load( + weight_ptr + checkpoint_row[:, None] * 320 + offsets[None, :], + mask=hidden_mask[:, None] & k_mask[None, :], + other=0.0, + eviction_policy="evict_first", + ) + gate = tl.sum(lora[None, :] * weight.to(tl.float32), axis=1) + # The communication kernel applies the original FP16 gate boundary, + # sigmoid, rank-ordered FP32 FMA, and final FP16 materialization. + tl.store(gate_ptr + hidden, gate, mask=hidden_mask) + + +@triton.jit +def _qwen38_hc_up_hidden_shard_kernel( + lora_ptr, + weight_ptr, + branches_ptr, + out_ptr, + TP_RANK: tl.constexpr, +): + """Mix all four branches locally for two of this rank's 640 hidden rows.""" + rows = tl.arange(0, 8) + hidden = tl.program_id(0) * 2 + rows // 4 + checkpoint_row = (rows % 4) * 2560 + TP_RANK * 640 + hidden + offsets = tl.arange(0, 512) + lora = tl.load(lora_ptr + offsets, offsets < 320, 0).to(tl.float32) + weight = tl.load( + weight_ptr + checkpoint_row[:, None] * 320 + offsets[None, :], + offsets[None, :] < 320, + 0, + ) + # Keep the existing two-K-warp reduction, FP16 gate boundary, and + # branch-ordered FP32 FMA. Only row ownership changes; weights are neither + # repacked nor duplicated, and prefill keeps its original layout. + gate = tl.sum(lora[None, :] * weight.to(tl.float32), axis=1) + gate = gate.to(tl.float16).to(tl.float32).reshape((2, 4)) + branches = tl.load(branches_ptr + checkpoint_row).to(tl.float32).reshape((2, 4)) + result = tl.full((2,), 0, tl.float32) + for branch in tl.static_range(4): + index = tl.full((2, 1), branch, tl.int32) + g = tl.gather(gate, index, 1).reshape((2,)) + x = tl.gather(branches, index, 1).reshape((2,)) + result = tl.fma(tl.sigmoid(g), x, result) + tl.store(out_ptr + tl.program_id(0) * 2 + tl.arange(0, 2), result / 4) + + @triton.jit def _qwen38_hc_up_gate_mix_kernel( lora_ptr, @@ -101,6 +206,52 @@ def _qwen38_hc_up_gate_mix_kernel( tl.store(out_ptr + hidden, result / HC_COUNT) +@triton.jit +def _qwen38_hc_up_gate_mix_row4_kernel( + lora_ptr, + weight_ptr, + x_ptr, + out_ptr, + K: tl.constexpr, + HC_DIMENSION: tl.constexpr, + HC_COUNT: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Reuse the low-rank input across four bitwise-equivalent output rows.""" + hidden = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) + offsets = tl.arange(0, BLOCK_K) + hidden_mask = hidden < HC_DIMENSION + k_mask = offsets < K + lora = tl.load( + lora_ptr + offsets, + mask=k_mask, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + + result = tl.zeros((BLOCK_N,), dtype=tl.float32) + for stream in tl.static_range(HC_COUNT): + row = stream * HC_DIMENSION + hidden + weight = tl.load( + weight_ptr + row[:, None] * K + offsets[None, :], + mask=hidden_mask[:, None] & k_mask[None, :], + other=0.0, + eviction_policy="evict_first", + ) + # Keep the established FP32 reduction and FP16 gate boundary. Row + # tiling changes only work assignment and shares the lora read. + gate = tl.sum(lora[None, :] * weight.to(tl.float32), axis=1) + gate = gate.to(tl.float16).to(tl.float32) + branch = tl.load( + x_ptr + stream * HC_DIMENSION + hidden, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + result += tl.sigmoid(gate) * branch + tl.store(out_ptr + hidden, result / HC_COUNT, mask=hidden_mask) + + def _runtime_ok( x: torch.Tensor, down_weight: torch.Tensor, up_weight: torch.Tensor ) -> bool: @@ -140,6 +291,68 @@ def _qwen38_sm70_fp16_fused_hc( gate = torch.nn.functional.linear(lora, up_weight) block = torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, _HC_COUNT) return block, injection + try: + from vllm.distributed.parallel_state import get_tp_group + + device_communicator = get_tp_group().device_communicator + custom_ar = getattr(device_communicator, "ca_comm", None) + except (AssertionError, AttributeError, RuntimeError, ValueError): + custom_ar = None + + if custom_ar is not None and custom_ar.can_sm70_qwen38_hc_shard(x): + tp_rank = int(custom_ar.rank) + local_down = x.new_empty((1, 88)) + gathered_down = x.new_empty((1, 336)) + block = x.new_empty((1, _HC_DIM)) + _qwen38_hc_down_local_shard_kernel[(88,)]( + x, + down_weight, + local_down, + TP_RANK=tp_rank, + num_warps=4, + ) + custom_ar.sm70_qwen38_hc_down_allgather(local_down, gathered_down) + if custom_ar.supports_sm70_qwen38_hc_up_mix_allgather(): + custom_ar.sm70_qwen38_hc_up_mix_allgather( + gathered_down, up_weight, x, block + ) + logger.info_once( + "SM70 Qwen3.8 exact TP4 fused FP16 HC up/mix/gather enabled." + ) + return block, gathered_down[..., _HC_RANK : _HC_RANK + _HC_COUNT] + if custom_ar.supports_sm70_qwen38_hc_output_allgather(): + local_block = x.new_empty((1, _HC_DIM // _HC_COUNT)) + _qwen38_hc_up_hidden_shard_kernel[(320,)]( + gathered_down, + up_weight, + x, + local_block, + TP_RANK=tp_rank, + num_warps=8, + ) + custom_ar.sm70_qwen38_hc_output_allgather(local_block, block) + logger.info_once( + "SM70 Qwen3.8 exact TP4 hidden-sharded FP16 HC route enabled." + ) + return block, gathered_down[..., _HC_RANK : _HC_RANK + _HC_COUNT] + + # An older wheel/sidecar can still use the established gate-sharded + # route. Never pass its opaque communicator to a different DSO. + local_gate = x.new_empty((1, _HC_DIM)) + _qwen38_hc_up_local_gate_kernel[(triton.cdiv(_HC_DIM, 8),)]( + gathered_down, + up_weight, + local_gate, + TP_RANK=tp_rank, + BLOCK_N=8, + num_warps=8, + ) + custom_ar.sm70_qwen38_hc_gate_mix(local_gate, x, block) + logger.info_once( + "SM70 Qwen3.8 exact TP4-sharded checkpoint-FP16 HC route enabled." + ) + return block, gathered_down[..., _HC_RANK : _HC_RANK + _HC_COUNT] + lora = x.new_empty((1, _HC_RANK)) injection = x.new_empty((1, _HC_COUNT)) block = x.new_empty((1, _HC_DIM)) @@ -154,7 +367,7 @@ def _qwen38_sm70_fp16_fused_hc( HC_COUNT=_HC_COUNT, num_warps=4, ) - _qwen38_hc_up_gate_mix_kernel[(_HC_DIM,)]( + _qwen38_hc_up_gate_mix_row4_kernel[(triton.cdiv(_HC_DIM, 4),)]( lora, up_weight, x, @@ -162,8 +375,9 @@ def _qwen38_sm70_fp16_fused_hc( K=_HC_RANK, HC_DIMENSION=_HC_DIM, HC_COUNT=_HC_COUNT, + BLOCK_N=4, BLOCK_K=512, - num_warps=2, + num_warps=8, ) logger.info_once("SM70 Qwen3.8 fused checkpoint-FP16 HC M=1 route enabled.") return block, injection @@ -243,6 +457,8 @@ def enable_qwen38_sm70_fp16_fused_hc( __all__ = [ + "_qwen38_hc_down_local_shard_kernel", + "_qwen38_hc_up_local_gate_kernel", "enable_qwen38_sm70_fp16_fused_hc", "maybe_apply_qwen38_sm70_fp16_fused_hc", ]