From 1c2d74a8dba969df509c9220f1042f17e55da656 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 01:57:37 +0800 Subject: [PATCH 01/22] [Perf][SM70] Record Qwen3.8 decode trace baseline Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_nvfp4_decode.md | 51 +++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index df8b0554d2..e2f40c94f8 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -816,3 +816,54 @@ 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. From efdf85a7049523bce930872bfdddf8653c3a1622 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 03:43:33 +0800 Subject: [PATCH 02/22] [Kernel][SM70] Fuse exact Qwen3.8 decode epilogues Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 5 + csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu | 142 ++++++++++++++++++ csrc/torch_bindings.cpp | 7 + tests/models/qwen4_exp/test_sm70_fp16_gemv.py | 49 ++++++ .../test_sm70_modelopt_mixed_nvfp4.py | 13 ++ tests/quantization/test_sm70_online_qpn8.py | 11 ++ vllm/_sm70_ops.py | 47 ++++++ vllm/envs.py | 6 + .../layers/quantization/nvfp4_sm70_moe.py | 34 +++++ vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py | 51 ++++++- 10 files changed, 363 insertions(+), 2 deletions(-) diff --git a/csrc/ops.h b/csrc/ops.h index 57d4bfb1c6..8aa74d6cc5 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -523,6 +523,11 @@ 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_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/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu index ad00901b59..38c3cc17e0 100644 --- a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -243,6 +243,105 @@ __global__ void nvfp4_qpn_m1_sm70_kernel(const half* __restrict__ input, } } +// 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); + } +} + template void launch_mxfp4_qpn_m1(torch::Tensor out, torch::Tensor input, torch::Tensor weights, torch::Tensor scales, @@ -406,6 +505,49 @@ 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_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..5e0d811453 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -697,6 +697,13 @@ 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); + // 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( diff --git a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py index 9a4bfe1e3b..0a944f1b3c 100644 --- a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py +++ b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py @@ -2,9 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch import vllm.envs as envs from vllm.models.qwen4_exp.nvidia.sm70_fp16_gemv import _plan_for +from vllm.models.qwen4_exp.nvidia.sm70_fp16_hc import ( + _qwen38_hc_up_gate_mix_kernel, + _qwen38_hc_up_gate_mix_row4_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 +76,45 @@ 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.cuda.synchronize() + assert torch.equal(actual, reference) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 4b50a35d42..a093e199bf 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -119,6 +119,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..ac380772a2 100644 --- a/tests/quantization/test_sm70_online_qpn8.py +++ b/tests/quantization/test_sm70_online_qpn8.py @@ -77,6 +77,17 @@ 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() + + @pytest.mark.parametrize( ("prefix", "k", "n", "expected"), [ diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 6fc4440ca1..bda6c22a41 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -154,6 +154,12 @@ 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_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 +1548,47 @@ 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_moe_qpn_mtp5_sm70_out( out: torch.Tensor, input: torch.Tensor, diff --git a/vllm/envs.py b/vllm/envs.py index d50e8dd634..520bb4868b 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 @@ -1903,6 +1904,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" ), diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 10915381e9..098d618d39 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -431,6 +431,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() @@ -646,6 +665,9 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: ) layer.sm70_nvfp4_qwen38_fused_swiglu_prefill = fused_swiglu_prefill 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) @@ -959,6 +981,18 @@ def apply( 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/models/qwen4_exp/nvidia/sm70_fp16_hc.py b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py index 4574bf3d6f..03f3d2c4eb 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py +++ b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py @@ -101,6 +101,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: @@ -154,7 +200,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 +208,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 From 5f2fbede02ce09ddec52549e9b8d4494b5c24721 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:02:08 +0800 Subject: [PATCH 03/22] [Kernel][SM70] Push Qwen3.8 M1 sum2 collectives Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/custom_all_reduce.cuh | 8 +++++++- tests/v1/spec_decode/test_dflash2.py | 13 +++++++++++++ .../device_communicators/custom_all_reduce.py | 7 ++++++- vllm/envs.py | 9 +++++++++ 4 files changed, 35 insertions(+), 2 deletions(-) diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 66e6765ff3..1c06021da9 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -1862,10 +1862,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/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/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 3bf6f6cf31..6cbbbb52e4 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, ) diff --git a/vllm/envs.py b/vllm/envs.py index 520bb4868b..f277420af5 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -233,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 @@ -2120,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"), From 355c67eb554358e6a2e55886ce008658a05b5cef Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:08:42 +0800 Subject: [PATCH 04/22] [CI][SM70] Format exact Qwen3.8 kernels Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 9 ++++---- csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu | 21 +++++++++---------- tests/models/qwen4_exp/test_sm70_fp16_gemv.py | 2 +- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/csrc/ops.h b/csrc/ops.h index 8aa74d6cc5..fc0f92df1a 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -523,10 +523,11 @@ 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_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_moe_qpn_mtp5_sm70_out(torch::Tensor out, torch::Tensor input, torch::Tensor weights, torch::Tensor scales, diff --git a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu index 38c3cc17e0..466a9ceddc 100644 --- a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -282,15 +282,14 @@ __global__ void nvfp4_qwen38_w2_direct_reduce_kernel( #pragma unroll for (int group = 0; group < kGroupsK16; ++group) { const size_t tile_group_base = - (static_cast(tile) * kGroupsK8 + group * 2) * 32 + - packed_col; + (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)); + 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); @@ -316,8 +315,7 @@ __global__ void nvfp4_qwen38_w2_direct_reduce_kernel( #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); + const int local_col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); route_outputs[route][quadpair * 8 + local_col] = __float2half(accum[index]); } @@ -505,10 +503,11 @@ 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) { +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(), @@ -539,7 +538,7 @@ void nvfp4_qwen38_w2_direct_reduce_out( const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); nvfp4_qwen38_w2_direct_reduce_kernel<<<80, 320, 0, - at::cuda::getCurrentCUDAStream()>>>( + at::cuda::getCurrentCUDAStream()>>>( reinterpret_cast(input.data_ptr()), reinterpret_cast(weights.data_ptr()), reinterpret_cast(scales.data_ptr()), diff --git a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py index 0a944f1b3c..a5a15cce45 100644 --- a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py +++ b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py @@ -116,5 +116,5 @@ def test_qwen38_sm70_hc_up_row4_is_bitwise() -> None: BLOCK_K=512, num_warps=8, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert torch.equal(actual, reference) From be65b288ae81848974eae0d0a9d55d1d9bbea33a Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:11:51 +0800 Subject: [PATCH 05/22] [Doc][SM70] Record exact Qwen3.8 token gains Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_nvfp4_decode.md | 37 +++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index e2f40c94f8..8121fde0e0 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -867,3 +867,40 @@ 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. + +The isolated savings sum to `0.498 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 after another material exact candidate lands. + +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. The retained HC changes do not quantize FP16 tensors or relax any +quality gate. From 19c21022b4ea00903c13bcb80820ec6bb4217320 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:50:46 +0800 Subject: [PATCH 06/22] [Kernel][SM70] Compact exact QSA decode selection Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/qsa_lexicographic_topk.cuh | 255 +++++++++++++++++++++++- docs/design/sm70_qwen38_nvfp4_decode.md | 23 ++- tests/kernels/test_top_k_per_row.py | 24 ++- 3 files changed, 293 insertions(+), 9 deletions(-) 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/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index 8121fde0e0..eafc2fd57d 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -889,10 +889,19 @@ boundaries: 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. - -The isolated savings sum to `0.498 ms/token`; they are not an end-to-end TPOT +- 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. + +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 after another material exact candidate lands. +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 @@ -902,5 +911,9 @@ 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. The retained HC changes do not quantize FP16 tensors or relax any -quality gate. +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. 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) From e4d46d97c25a162eaa35a157721604e83174f3e5 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 04:59:26 +0800 Subject: [PATCH 07/22] [Core][SM70] Allow exact QSA source fragment Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_nvfp4_decode.md | 4 +++- tests/models/qwen4_exp/test_qsa_ops.py | 19 +++++++++++++++++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 28 ++++++++++++++++++++++++- 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index eafc2fd57d..568e7d72e6 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -897,7 +897,9 @@ boundaries: `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. + 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 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/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index b8b370ba18..a5015ac2d3 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 @@ -1072,6 +1089,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 +1358,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, From 8298d3827903f793e4186b1362b397af6c2454df Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 05:14:52 +0800 Subject: [PATCH 08/22] [Doc][SM70] Close exact HC schedule screens Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_nvfp4_decode.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index 568e7d72e6..8a3a90dae9 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -919,3 +919,14 @@ 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. From 728b5016b456adfe17cf8f8334987fabc2e0ef86 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 06:17:41 +0800 Subject: [PATCH 09/22] [Kernel][SM70] Shard exact Qwen3.8 HC compute across TP4 Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/custom_all_reduce.cu | 316 ++++++++++++++++++ csrc/ops.h | 4 + csrc/torch_bindings.cpp | 9 + docs/design/sm70_qwen38_nvfp4_decode.md | 37 ++ tests/models/qwen4_exp/test_sm70_fp16_gemv.py | 78 +++++ vllm/_custom_ops.py | 17 + .../device_communicators/custom_all_reduce.py | 25 ++ vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py | 109 ++++++ 8 files changed, 595 insertions(+) diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index e55f515ebb..a32aa343ba 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -87,6 +87,237 @@ 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; + +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 = kSm70Tp4PushAllreduceBytes / sizeof(P); + + auto* local_storage = + const_cast(reinterpret_cast(push_buffers.ptrs[Rank])); + auto* local_epochs = reinterpret_cast(local_storage); + const uint32_t epoch = local_epochs[0]; + 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 + kSm70Tp4PushAllreduceSignalBytes + + (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 + kSm70Tp4PushAllreduceSignalBytes + + (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 + kSm70Tp4PushAllreduceSignalBytes + + (epoch_offset + source_rank * kPackedStride) * sizeof(P); + sm70_push_store_volatile_16b(empty, source, offset); + } + } + + __syncthreads(); + if (threadIdx.x == 0) { + local_epochs[0] = (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; +} + +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 = kSm70Tp4PushAllreduceBytes / sizeof(P); + + auto* local_storage = + const_cast(reinterpret_cast(push_buffers.ptrs[Rank])); + auto* local_epochs = reinterpret_cast(local_storage); + const uint32_t epoch = local_epochs[blockIdx.x]; + 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 + kSm70Tp4PushAllreduceSignalBytes + + (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 + kSm70Tp4PushAllreduceSignalBytes + + (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 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 + kSm70Tp4PushAllreduceSignalBytes + + (epoch_offset + source_rank * kPackedStride) * sizeof(P); + sm70_push_store_volatile_16b(empty, source, offset); + } + } + + __syncthreads(); + if (threadIdx.x == 0) { + local_epochs[blockIdx.x] = (epoch + 1) % kSm70Tp4PushAllreduceEpochs; + } +} + +} // namespace vllm +#endif + /** * Performs an out-of-place allreduce and stores result in out. * @@ -413,6 +644,91 @@ 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 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/ops.h b/csrc/ops.h index fc0f92df1a..ab4d906af9 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -650,6 +650,10 @@ void sm70_tp4_reduce_scatter_gemma_rms_norm_all_gather( fptr_t reg_output_buffer, int64_t reg_buffer_sz_bytes, double epsilon); void all_reduce_sum2(fptr_t _fa, torch::Tensor& inp_a, torch::Tensor& inp_b, torch::Tensor& out); +void sm70_qwen38_hc_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 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/torch_bindings.cpp b/csrc/torch_bindings.cpp index 5e0d811453..2c0409647b 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -910,6 +910,15 @@ 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( "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 8a3a90dae9..e5a737e5c8 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -930,3 +930,40 @@ 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. diff --git a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py index a5a15cce45..09c50c23d4 100644 --- a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py +++ b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py @@ -5,10 +5,14 @@ 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_local_gate_kernel, ) from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON @@ -118,3 +122,77 @@ def test_qwen38_sm70_hc_up_row4_is_bitwise() -> None: ) 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) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index f2bea0da4b..7dca310a67 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3085,6 +3085,23 @@ 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 top1_argmax( fa: int, input_pair: torch.Tensor, diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 6cbbbb52e4..36974dce32 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -440,6 +440,31 @@ 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 sm70_tp2_all_reduce_gemma_rms_norm( self, inp: torch.Tensor, diff --git a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py index 03f3d2c4eb..add14d25c2 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py +++ b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py @@ -63,6 +63,77 @@ 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_gate_mix_kernel( lora_ptr, @@ -186,6 +257,42 @@ 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)) + local_gate = x.new_empty((1, _HC_DIM)) + 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) + _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)) @@ -290,6 +397,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", ] From aa1db9c29b787e61227876824a5abf0ca26152d1 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:34:20 +0800 Subject: [PATCH 10/22] [Bugfix][SM70] Isolate Qwen3.8 HC push channels Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/custom_all_reduce.cu | 33 ++++++++++++++++++------------- csrc/custom_all_reduce.cuh | 40 ++++++++++++++++++++++++++++++++++---- 2 files changed, 55 insertions(+), 18 deletions(-) diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index a32aa343ba..3439bdbb71 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -110,12 +110,14 @@ __global__ void __launch_bounds__(128, 1) using P = typename packed_t::P; constexpr int kElementsPerPack = P::size; constexpr int kPackedElements = kQwen38HcDownLocalElements / kElementsPerPack; - constexpr int kPackedStride = kSm70Tp4PushAllreduceBytes / sizeof(P); + 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); - const uint32_t epoch = local_epochs[0]; + 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; @@ -132,7 +134,7 @@ __global__ void __launch_bounds__(128, 1) if (destination_rank == Rank) continue; auto* destination_base = const_cast( reinterpret_cast(push_buffers.ptrs[destination_rank])); - void* destination = destination_base + kSm70Tp4PushAllreduceSignalBytes + + void* destination = destination_base + kSm70Qwen38HcDownPushOffset + (epoch_offset + Rank * kPackedStride) * sizeof(P); sm70_push_store_volatile_16b(value, destination, offset); } @@ -145,7 +147,7 @@ __global__ void __launch_bounds__(128, 1) for (int source_rank = 0; source_rank < ngpus; ++source_rank) { if (source_rank == Rank) continue; const void* source = - local_storage + kSm70Tp4PushAllreduceSignalBytes + + local_storage + kSm70Qwen38HcDownPushOffset + (epoch_offset + source_rank * kPackedStride) * sizeof(P); sm70_push_load_volatile_16b(peer_values[source_rank], source, offset); #pragma unroll @@ -189,7 +191,7 @@ __global__ void __launch_bounds__(128, 1) #pragma unroll for (int source_rank = 0; source_rank < ngpus; ++source_rank) { if (source_rank == Rank) continue; - void* source = local_storage + kSm70Tp4PushAllreduceSignalBytes + + void* source = local_storage + kSm70Qwen38HcDownPushOffset + (epoch_offset + source_rank * kPackedStride) * sizeof(P); sm70_push_store_volatile_16b(empty, source, offset); } @@ -197,7 +199,8 @@ __global__ void __launch_bounds__(128, 1) __syncthreads(); if (threadIdx.x == 0) { - local_epochs[0] = (epoch + 1) % kSm70Tp4PushAllreduceEpochs; + local_epochs[kSm70Qwen38HcDownEpochIndex] = + (epoch + 1) % kSm70Tp4PushAllreduceEpochs; } } @@ -233,12 +236,14 @@ __global__ void __launch_bounds__(512, 1) int packed_elements) { static_assert(ngpus == kSm70Tp4PushAllreduceWorldSize); using P = typename packed_t::P; - constexpr int kPackedStride = kSm70Tp4PushAllreduceBytes / sizeof(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); - const uint32_t epoch = local_epochs[blockIdx.x]; + 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; @@ -255,7 +260,7 @@ __global__ void __launch_bounds__(512, 1) if (destination_rank == Rank) continue; auto* destination_base = const_cast( reinterpret_cast(push_buffers.ptrs[destination_rank])); - void* destination = destination_base + kSm70Tp4PushAllreduceSignalBytes + + void* destination = destination_base + kSm70Qwen38HcGatePushOffset + (epoch_offset + Rank * kPackedStride) * sizeof(P); sm70_push_store_volatile_16b(value, destination, offset); } @@ -268,7 +273,7 @@ __global__ void __launch_bounds__(512, 1) for (int source_rank = 0; source_rank < ngpus; ++source_rank) { if (source_rank == Rank) continue; const void* source = - local_storage + kSm70Tp4PushAllreduceSignalBytes + + local_storage + kSm70Qwen38HcGatePushOffset + (epoch_offset + source_rank * kPackedStride) * sizeof(P); sm70_push_load_volatile_16b(peer_values[source_rank], source, offset); #pragma unroll @@ -303,7 +308,7 @@ __global__ void __launch_bounds__(512, 1) #pragma unroll for (int source_rank = 0; source_rank < ngpus; ++source_rank) { if (source_rank == Rank) continue; - void* source = local_storage + kSm70Tp4PushAllreduceSignalBytes + + void* source = local_storage + kSm70Qwen38HcGatePushOffset + (epoch_offset + source_rank * kPackedStride) * sizeof(P); sm70_push_store_volatile_16b(empty, source, offset); } @@ -311,7 +316,7 @@ __global__ void __launch_bounds__(512, 1) __syncthreads(); if (threadIdx.x == 0) { - local_epochs[blockIdx.x] = (epoch + 1) % kSm70Tp4PushAllreduceEpochs; + local_epochs[epoch_index] = (epoch + 1) % kSm70Tp4PushAllreduceEpochs; } } diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 1c06021da9..2c14ceeaf7 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -75,10 +75,34 @@ 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 kSm70Tp4PushAllreduceBufferBytes = + kSm70Qwen38HcGatePushOffset + kSm70Tp4PushAllreduceEpochs * + kSm70Tp4PushAllreduceWorldSize * + kSm70Qwen38HcGatePushBytes; +static_assert(kSm70Qwen38HcGateEpochIndexBase + kSm70Qwen38HcGatePushBlocks <= + kSm70Qwen38HcPushSignalBytes / sizeof(uint32_t)); inline int sm70_tp4_push_allreduce_blocks(size_t bytes) { if (bytes == kSm70Tp4PushAllreduceBytes) { @@ -1505,11 +1529,19 @@ 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( - local_data, kSm70Tp4PushAllreduceSentinelByte, - kSm70Tp4PushAllreduceBufferBytes - kSm70Tp4PushAllreduceSignalBytes)); + hc_data, kSm70Tp4PushAllreduceSentinelByte, + kSm70Tp4PushAllreduceBufferBytes - kSm70Qwen38HcDownPushOffset)); sm70_tp4_push_buffers_registered_ = true; } From 234d6bea91a89352150d3aaf14b52495e47fac02 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:34:32 +0800 Subject: [PATCH 11/22] [Kernel][SM70] Tune Qwen3.8 NVFP4 W13 split plan Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/quantization/test_sm70_modelopt_mixed_nvfp4.py | 5 +++++ vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py | 5 ++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index a093e199bf..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)], diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 098d618d39..6a8d04b92a 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 From 321fa105b622e9500a33941f7b3146cd2f51375e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:00:50 +0800 Subject: [PATCH 12/22] [Kernel][SM70] Fuse exact Qwen3.8 PLE decode primitives Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 39 ++++ vllm/models/qwen4_exp/nvidia/ple_layer.py | 239 ++++++++++++++++++++- 2 files changed, 276 insertions(+), 2 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index ce84e94a98..7d12848005 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -44993,3 +44993,42 @@ 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. 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 ) From b9e1c0146626ea6d7a40cc65765027974cc66a80 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:50:42 +0800 Subject: [PATCH 13/22] [Kernel][SM70] Fuse exact Qwen3.8 W13 SwiGLU decode Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 5 ++ csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu | 81 ++++++++++++++++++- csrc/torch_bindings.cpp | 7 ++ docs/design/sm70_v100_migration_control.md | 17 ++++ vllm/_sm70_ops.py | 44 ++++++++++ .../layers/quantization/nvfp4_sm70_moe.py | 66 +++++++++++---- 6 files changed, 202 insertions(+), 18 deletions(-) diff --git a/csrc/ops.h b/csrc/ops.h index ab4d906af9..82686c260e 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -529,6 +529,11 @@ void nvfp4_qwen38_w2_direct_reduce_out(torch::Tensor out, torch::Tensor input, 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, diff --git a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu index 466a9ceddc..ca75bfc9a3 100644 --- a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu +++ b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu @@ -138,7 +138,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 +153,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 +243,23 @@ __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; + } } } @@ -371,6 +391,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, @@ -547,6 +586,40 @@ void nvfp4_qwen38_w2_direct_reduce_out(torch::Tensor out, torch::Tensor input, 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 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 2c0409647b..df70748ebe 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -704,6 +704,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { 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( diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 7d12848005..b67db2028e 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45032,3 +45032,20 @@ Interpretation: 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. diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index bda6c22a41..e085347054 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -160,6 +160,12 @@ def has_nvfp4_qwen38_w2_direct_reduce() -> bool: ) +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_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( @@ -1589,6 +1595,44 @@ def _nvfp4_qwen38_w2_direct_reduce_out_sidecar_fake( 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 nvfp4_moe_qpn_mtp5_sm70_out( out: torch.Tensor, input: torch.Tensor, diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 6a8d04b92a..13592cf777 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -550,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 ) @@ -667,6 +681,7 @@ 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 @@ -699,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 " @@ -969,21 +989,39 @@ 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) ): From 396d432076df5d0f1e8bea2a165e9006975abdc1 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:13:06 +0800 Subject: [PATCH 14/22] [Kernel][SM70] Fuse exact Qwen3.8 shared gate Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 3 + csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu | 80 ++++++++++++++++++++ csrc/torch_bindings.cpp | 6 ++ docs/design/sm70_v100_migration_control.md | 22 ++++++ tests/quantization/test_sm70_online_qpn8.py | 11 +++ vllm/_sm70_ops.py | 36 +++++++++ vllm/model_executor/models/qwen2_moe.py | 38 ++++++---- 7 files changed, 181 insertions(+), 15 deletions(-) diff --git a/csrc/ops.h b/csrc/ops.h index 82686c260e..ae40357150 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); diff --git a/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/mxfp4_qpn_m1_sm70.cu index ca75bfc9a3..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; @@ -620,6 +672,34 @@ void nvfp4_qwen38_w13_fused_swiglu_out(torch::Tensor out, torch::Tensor input, 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 df70748ebe..8779503fae 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); diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index b67db2028e..653024ed5d 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45049,3 +45049,25 @@ Interpretation: - 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. diff --git a/tests/quantization/test_sm70_online_qpn8.py b/tests/quantization/test_sm70_online_qpn8.py index ac380772a2..e5090c3c19 100644 --- a/tests/quantization/test_sm70_online_qpn8.py +++ b/tests/quantization/test_sm70_online_qpn8.py @@ -88,6 +88,17 @@ def test_nvfp4_w2_direct_reduce_capability_is_explicit(monkeypatch): 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/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index e085347054..fe85479fb7 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -166,6 +166,12 @@ def has_nvfp4_qwen38_w13_fused_swiglu() -> bool: ) +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( @@ -1633,6 +1639,36 @@ def _nvfp4_qwen38_w13_fused_swiglu_out_sidecar_fake( 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/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 From 30f81105621e9f39e6b3bf9d816f77d63acd8307 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:33:17 +0800 Subject: [PATCH 15/22] [Kernel][SM70] Fuse exact Qwen3.8 QSA output gate Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 15 ++++ tests/models/qwen4_exp/test_qsa_reference.py | 18 ++++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 86 ++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/qsa.py | 14 +++- 4 files changed, 130 insertions(+), 3 deletions(-) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 653024ed5d..27022380a3 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45071,3 +45071,18 @@ Interpretation: 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. 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/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index a5015ac2d3..8c4e36c1dd 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -540,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, @@ -552,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, @@ -677,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 @@ -718,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, @@ -757,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, @@ -1943,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", @@ -1996,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 @@ -2025,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] @@ -2078,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), @@ -2090,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], @@ -2116,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/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( From aaf63696b6e00936153e70a3d6daaa4e77a30593 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:12:07 +0800 Subject: [PATCH 16/22] [Kernel][SM70] Shard exact HC mixing by hidden coordinate Preserve FP32 reductions and FP16 boundaries while gathering 640 mixed hidden values per rank. Keep older communicator DSOs on the gate-sharded route. Record real-weight screens and reject slower publication/split prototypes; add reproducible production and auxiliary-stream gates. Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/kernels/benchmark_sm70_hc_tp4.py | 225 ++++++++++++++++++ csrc/custom_all_reduce.cu | 81 ++++++- csrc/ops.h | 2 + csrc/torch_bindings.cpp | 5 + docs/design/sm70_qwen38_nvfp4_decode.md | 59 +++++ docs/design/sm70_v100_migration_control.md | 60 +++++ .../test_custom_all_reduce_dispatch.py | 37 +++ tests/models/qwen4_exp/test_sm70_fp16_gemv.py | 47 ++++ vllm/_custom_ops.py | 19 ++ .../device_communicators/custom_all_reduce.py | 8 + vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py | 54 ++++- 11 files changed, 586 insertions(+), 11 deletions(-) create mode 100644 benchmarks/kernels/benchmark_sm70_hc_tp4.py diff --git a/benchmarks/kernels/benchmark_sm70_hc_tp4.py b/benchmarks/kernels/benchmark_sm70_hc_tp4.py new file mode 100644 index 0000000000..87189460be --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_hc_tp4.py @@ -0,0 +1,225 @@ +# 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 + + 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 3439bdbb71..f079a82a18 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -100,6 +100,8 @@ constexpr int kQwen38HcDownGatheredElements = constexpr int kQwen38HcGateLocalElements = 2560; constexpr int kQwen38HcGateGatheredElements = kQwen38HcGateLocalElements * kSm70Tp4PushAllreduceWorldSize; +constexpr int kQwen38HcOutputLocalElements = + kQwen38HcGateLocalElements / kSm70Tp4PushAllreduceWorldSize; template __global__ void __launch_bounds__(128, 1) @@ -227,7 +229,7 @@ DINLINE float qwen38_hc_divide_by_count(float value) { return result; } -template +template __global__ void __launch_bounds__(512, 1) sm70_qwen38_hc_gate_push_mix(RankData push_buffers, const half* __restrict__ local_gate, @@ -285,18 +287,31 @@ __global__ void __launch_bounds__(512, 1) if (!has_empty_slot) break; } - #pragma unroll - for (int element = 0; element < P::size; ++element) { - const int hidden = offset * P::size + element; - float result = 0.0f; + 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) { - 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); + 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)); } - output[hidden] = __float2half_rn(qwen38_hc_divide_by_count(result)); } P empty; @@ -734,6 +749,52 @@ void sm70_qwen38_hc_gate_mix(fptr_t _fa, torch::Tensor& local_gate, #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/ops.h b/csrc/ops.h index ae40357150..314a5a653b 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -662,6 +662,8 @@ 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 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/torch_bindings.cpp b/csrc/torch_bindings.cpp index 8779503fae..040ec46066 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -932,6 +932,11 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { "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( "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 e5a737e5c8..a64e21a762 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -967,3 +967,62 @@ 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 retained improvement is about 2.3%, not the initial 20% screening target +and not an end-to-end tokens/s claim. The 100-tok/s endpoint 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. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 27022380a3..8bf21b7b09 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45086,3 +45086,63 @@ Interpretation: 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. Production GPU validation and the next combined full-model + quality/performance gate remain pending; do not claim 100 tok/s or promote + an endpoint from the isolated 0.041-ms saving. diff --git a/tests/distributed/test_custom_all_reduce_dispatch.py b/tests/distributed/test_custom_all_reduce_dispatch.py index 3703931fd3..599c3e2cd8 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 Mock + import pytest import torch +import vllm._custom_ops as ops from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce @@ -33,3 +37,36 @@ 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]) +def test_hc_output_gather_stays_in_communicator_dso( + monkeypatch: pytest.MonkeyPatch, sidecar_owner: bool, owner_has_op: bool +) -> 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) + other.sm70_qwen38_hc_output_allgather = Mock() + if owner_has_op: + owner.sm70_qwen38_hc_output_allgather = Mock() + monkeypatch.setattr(torch.ops, "_C_custom_ar", base) + monkeypatch.setattr(torch.ops, "_C_custom_ar_flashnext", sidecar) + assert ops.supports_sm70_qwen38_hc_output_allgather() == owner_has_op + if owner_has_op: + local = torch.empty(640) + output = torch.empty(2560) + ops.sm70_qwen38_hc_output_allgather(123, local, output) + owner.sm70_qwen38_hc_output_allgather.assert_called_once_with( + 123, local, output + ) + 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): + ops.sm70_qwen38_hc_output_allgather( + 123, torch.empty(640), torch.empty(2560) + ) + other.sm70_qwen38_hc_output_allgather.assert_not_called() diff --git a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py index 09c50c23d4..8b33802cc7 100644 --- a/tests/models/qwen4_exp/test_sm70_fp16_gemv.py +++ b/tests/models/qwen4_exp/test_sm70_fp16_gemv.py @@ -12,6 +12,7 @@ _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 @@ -196,3 +197,49 @@ def test_qwen38_sm70_hc_tp4_compute_shards_are_bitwise() -> None: 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/vllm/_custom_ops.py b/vllm/_custom_ops.py index 7dca310a67..03cf815364 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3102,6 +3102,25 @@ def sm70_qwen38_hc_gate_mix( _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 top1_argmax( fa: int, input_pair: torch.Tensor, diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 36974dce32..4e8b50484d 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -465,6 +465,14 @@ def sm70_qwen38_hc_gate_mix( ) -> 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 sm70_tp2_all_reduce_gemma_rms_norm( self, inp: torch.Tensor, diff --git a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py index add14d25c2..b05740b8d7 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py +++ b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py @@ -134,6 +134,40 @@ def _qwen38_hc_up_local_gate_kernel( 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, @@ -269,7 +303,6 @@ def _qwen38_sm70_fp16_fused_hc( tp_rank = int(custom_ar.rank) local_down = x.new_empty((1, 88)) gathered_down = x.new_empty((1, 336)) - local_gate = x.new_empty((1, _HC_DIM)) block = x.new_empty((1, _HC_DIM)) _qwen38_hc_down_local_shard_kernel[(88,)]( x, @@ -279,6 +312,25 @@ def _qwen38_sm70_fp16_fused_hc( num_warps=4, ) custom_ar.sm70_qwen38_hc_down_allgather(local_down, gathered_down) + 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, From 50f9fbe3749ecd673fee1aef361afd8db98e914a Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:24:42 +0800 Subject: [PATCH 17/22] [Doc][SM70] Record production HC hidden-shard gates Record exact real-weight TP4 and auxiliary-stream CUDA Graph results, paired Mix-only timings, runtime/toolchain, and process cleanup. Keep full-model acceptance separate. Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_nvfp4_decode.md | 25 +++++++++++++++++++--- docs/design/sm70_v100_migration_control.md | 25 +++++++++++++++++++--- 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index a64e21a762..c5cb622324 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -1008,9 +1008,28 @@ 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 retained improvement is about 2.3%, not the initial 20% screening target -and not an end-to-end tokens/s claim. The 100-tok/s endpoint target remains -unproven by this change. +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: diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 8bf21b7b09..df4f0c6722 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45143,6 +45143,25 @@ Interpretation: 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. Production GPU validation and the next combined full-model - quality/performance gate remain pending; do not claim 100 tok/s or promote - an endpoint from the isolated 0.041-ms saving. + 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. From 96140649c9b26b8332335f760ec9612607838b4f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 5 Sep 2026 14:37:24 +0800 Subject: [PATCH 18/22] [Benchmark] Add complete HC latency and exactness gate Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_sm70_hc_full_chain.py | 263 ++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 81 ++++++ 2 files changed, 344 insertions(+) create mode 100644 benchmarks/kernels/benchmark_sm70_hc_full_chain.py 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..911a809078 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_hc_full_chain.py @@ -0,0 +1,263 @@ +# 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 + +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) + args = parser.parse_args() + 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") + 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)) + + 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(hidden: bool): + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + outputs = [] + with ( + patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp), + patch.object( + comm, + "supports_sm70_qwen38_hc_output_allgather", + return_value=hidden, + ), + comm.capture(), + torch.cuda.graph(graph), + ): + 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)) + outputs.extend(finish(state, injection)) + torch.cuda.synchronize() + dist.barrier() + return graph, outputs + + graphs = {"gate": capture(False), "hidden": capture(True)} + mismatches = 0 + 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) + for graph, _ in graphs.values(): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + expected = torch.cat( + [x.flatten().view(torch.int16) for x in graphs["gate"][1]] + ) + actual = torch.cat( + [x.flatten().view(torch.int16) for x in graphs["hidden"][1]] + ) + mismatches += int(torch.count_nonzero(expected != actual)) + quality = [None] * 4 + dist.all_gather_object( + quality, {"rank": rank, "mismatches": 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 graph, _ in graphs.values(): + for _ in range(args.warmup): + graph.replay() + torch.cuda.synchronize() + dist.barrier() + samples = {mode: [] for mode in graphs} + for repeat in range(3): + modes = list(graphs) if repeat % 2 == 0 else list(graphs)[::-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() + 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, + "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/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index df4f0c6722..8dd4274cba 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45165,3 +45165,84 @@ Interpretation: 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. From 8ac72cb0d312bcd747cdc06b601cb1ee880f4052 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:00:04 +0800 Subject: [PATCH 19/22] [Doc] Record exact HC fusion screens and GPU reservation guard Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 31 ++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 8dd4274cba..692deb4733 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45246,3 +45246,34 @@ Interpretation: 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. From c13547ab7ddbd992cd465104484898be6a81dd04 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:03:58 +0800 Subject: [PATCH 20/22] [Doc] Record exact complete HC vector fusion gain Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_v100_migration_control.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 692deb4733..2b8bce252d 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45277,3 +45277,20 @@ Interpretation: 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. From 0303b82d1ec8fd9549d75018995939bbee63846e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:21:13 +0800 Subject: [PATCH 21/22] [Kernel] Fuse exact SM70 HC up projection and output gather Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_sm70_hc_full_chain.py | 137 ++++++++++++--- benchmarks/kernels/benchmark_sm70_hc_tp4.py | 16 +- csrc/custom_all_reduce.cu | 166 ++++++++++++++++++ csrc/custom_all_reduce.cuh | 19 +- csrc/ops.h | 5 + csrc/torch_bindings.cpp | 5 + docs/design/sm70_qwen38_nvfp4_decode.md | 41 +++++ docs/design/sm70_v100_migration_control.md | 30 ++++ .../test_custom_all_reduce_dispatch.py | 75 ++++++-- vllm/_custom_ops.py | 16 ++ .../device_communicators/custom_all_reduce.py | 12 ++ vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py | 8 + 12 files changed, 487 insertions(+), 43 deletions(-) diff --git a/benchmarks/kernels/benchmark_sm70_hc_full_chain.py b/benchmarks/kernels/benchmark_sm70_hc_full_chain.py index 911a809078..2b29c0f784 100644 --- a/benchmarks/kernels/benchmark_sm70_hc_full_chain.py +++ b/benchmarks/kernels/benchmark_sm70_hc_full_chain.py @@ -27,6 +27,7 @@ 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 ( @@ -45,7 +46,22 @@ def main() -> None: 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") @@ -85,6 +101,8 @@ def ensure_exclusive() -> None: 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" @@ -115,6 +133,20 @@ def get(name: str) -> torch.Tensor: 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( @@ -128,21 +160,31 @@ def finish(state: torch.Tensor, injection: torch.Tensor): finish(initial, torch.zeros((1, 4), device="cuda", dtype=torch.float16)) torch.cuda.synchronize() - def capture(hidden: bool): + 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=hidden, + 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. @@ -159,13 +201,51 @@ def capture(hidden: bool): 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 + 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 - graphs = {"gate": capture(False), "hidden": capture(True)} - mismatches = 0 for case in range(args.quality_inputs): initial.normal_(generator=gen) cores.normal_(generator=gen) @@ -175,34 +255,35 @@ def capture(hidden: bool): elif case == 1: initial.mul_(0.01) cores.mul_(0.01) - for graph, _ in graphs.values(): - graph.replay() - torch.cuda.synchronize() - dist.barrier() - expected = torch.cat( - [x.flatten().view(torch.int16) for x in graphs["gate"][1]] - ) - actual = torch.cat( - [x.flatten().view(torch.int16) for x in graphs["hidden"][1]] - ) - mismatches += int(torch.count_nonzero(expected != actual)) + 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": mismatches}, group=group + 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 graph, _ in graphs.values(): + 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 graphs} + samples = {mode: [] for mode in timed_modes} for repeat in range(3): - modes = list(graphs) if repeat % 2 == 0 else list(graphs)[::-1] + modes = timed_modes if repeat % 2 == 0 else timed_modes[::-1] for mode in modes: ensure_exclusive() graph = graphs[mode][0] @@ -225,6 +306,18 @@ def capture(hidden: bool): ) 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( @@ -248,6 +341,10 @@ def capture(hidden: bool): "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()}, } diff --git a/benchmarks/kernels/benchmark_sm70_hc_tp4.py b/benchmarks/kernels/benchmark_sm70_hc_tp4.py index 87189460be..a08b08bec8 100644 --- a/benchmarks/kernels/benchmark_sm70_hc_tp4.py +++ b/benchmarks/kernels/benchmark_sm70_hc_tp4.py @@ -127,11 +127,17 @@ def capture(hidden: bool, overlap: bool = False): dist.barrier() return graph, outputs, sums - graphs = { - "control": capture(False), - "hidden": capture(True), - "hidden_aux": capture(True, overlap=True), - } + # 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) diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index f079a82a18..5e246e1180 100644 --- a/csrc/custom_all_reduce.cu +++ b/csrc/custom_all_reduce.cu @@ -229,6 +229,137 @@ DINLINE float qwen38_hc_divide_by_count(float value) { 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, @@ -749,6 +880,41 @@ void sm70_qwen38_hc_gate_mix(fptr_t _fa, torch::Tensor& local_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) diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 2c14ceeaf7..99d7dce2c2 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -97,10 +97,19 @@ constexpr size_t kSm70Qwen38HcGatePushOffset = kSm70Qwen38HcDownPushOffset + kSm70Tp4PushAllreduceEpochs * kSm70Tp4PushAllreduceWorldSize * kSm70Qwen38HcDownPushBytes; -constexpr size_t kSm70Tp4PushAllreduceBufferBytes = +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)); @@ -1541,7 +1550,13 @@ class CustomAllreduce { static_cast(ptrs[rank_]) + kSm70Qwen38HcDownPushOffset; CUDACHECK(cudaMemset( hc_data, kSm70Tp4PushAllreduceSentinelByte, - kSm70Tp4PushAllreduceBufferBytes - kSm70Qwen38HcDownPushOffset)); + kSm70Qwen38HcUpFusedEpochOffset - kSm70Qwen38HcDownPushOffset)); + // The first fused packet uses generation 1; zero is initially invalid. + auto* hc_up = + static_cast(ptrs[rank_]) + kSm70Qwen38HcUpFusedEpochOffset; + CUDACHECK(cudaMemset( + hc_up, 0, + kSm70Tp4PushAllreduceBufferBytes - kSm70Qwen38HcUpFusedEpochOffset)); sm70_tp4_push_buffers_registered_ = true; } diff --git a/csrc/ops.h b/csrc/ops.h index 314a5a653b..37cad87e7d 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -664,6 +664,11 @@ 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/torch_bindings.cpp b/csrc/torch_bindings.cpp index 040ec46066..e46afd5851 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -937,6 +937,11 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { "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 c5cb622324..a023b693f1 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -1045,3 +1045,44 @@ 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. Production +dispatch and auxiliary-stream gates remain pending at this source update. +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 2b8bce252d..dbf6f35d47 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45294,3 +45294,33 @@ Interpretation: 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. diff --git a/tests/distributed/test_custom_all_reduce_dispatch.py b/tests/distributed/test_custom_all_reduce_dispatch.py index 599c3e2cd8..9c5ee6ca52 100644 --- a/tests/distributed/test_custom_all_reduce_dispatch.py +++ b/tests/distributed/test_custom_all_reduce_dispatch.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock import pytest import torch @@ -41,32 +41,75 @@ def test_should_custom_ar_rejects_unsupported_dtype(dtype: torch.dtype) -> None: @pytest.mark.parametrize("sidecar_owner", [False, True]) @pytest.mark.parametrize("owner_has_op", [False, True]) -def test_hc_output_gather_stays_in_communicator_dso( - monkeypatch: pytest.MonkeyPatch, sidecar_owner: bool, owner_has_op: bool +@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) - other.sm70_qwen38_hc_output_allgather = Mock() + setattr(other, op_name, Mock()) if owner_has_op: - owner.sm70_qwen38_hc_output_allgather = Mock() + setattr(owner, op_name, Mock()) monkeypatch.setattr(torch.ops, "_C_custom_ar", base) monkeypatch.setattr(torch.ops, "_C_custom_ar_flashnext", sidecar) - assert ops.supports_sm70_qwen38_hc_output_allgather() == owner_has_op + assert getattr(ops, f"supports_{op_name}")() == owner_has_op + tensors = [torch.empty(16) for _ in range(num_inputs + 1)] if owner_has_op: - local = torch.empty(640) - output = torch.empty(2560) - ops.sm70_qwen38_hc_output_allgather(123, local, output) - owner.sm70_qwen38_hc_output_allgather.assert_called_once_with( - 123, local, output - ) + 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): - ops.sm70_qwen38_hc_output_allgather( - 123, torch.empty(640), torch.empty(2560) - ) - other.sm70_qwen38_hc_output_allgather.assert_not_called() + 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/vllm/_custom_ops.py b/vllm/_custom_ops.py index 03cf815364..f8e23d6c7e 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3121,6 +3121,22 @@ def sm70_qwen38_hc_output_allgather( _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/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 4e8b50484d..2c8d6ffdce 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -473,6 +473,18 @@ def sm70_qwen38_hc_output_allgather( ) -> 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/models/qwen4_exp/nvidia/sm70_fp16_hc.py b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py index b05740b8d7..b271461663 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py +++ b/vllm/models/qwen4_exp/nvidia/sm70_fp16_hc.py @@ -312,6 +312,14 @@ def _qwen38_sm70_fp16_fused_hc( 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,)]( From 205acfb4da0d038344165dbf3abd92b8cb978b21 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:32:03 +0800 Subject: [PATCH 22/22] [Doc] Record registered HC fusion quality and full-chain gain Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_nvfp4_decode.md | 9 +++++++-- docs/design/sm70_v100_migration_control.md | 21 +++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/docs/design/sm70_qwen38_nvfp4_decode.md b/docs/design/sm70_qwen38_nvfp4_decode.md index a023b693f1..f4c7d77c41 100644 --- a/docs/design/sm70_qwen38_nvfp4_decode.md +++ b/docs/design/sm70_qwen38_nvfp4_decode.md @@ -1064,8 +1064,13 @@ 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. Production -dispatch and auxiliary-stream gates remain pending at this source update. +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. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index dbf6f35d47..bde32d70cc 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45324,3 +45324,24 @@ Interpretation: 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.