diff --git a/benchmarks/kernels/benchmark_kimi_k3_attn_res.py b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py new file mode 100644 index 0000000000..c8cf6f9acc --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_attn_res.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark vLLM Kimi-K3 AttnRes kernels on NVIDIA Blackwell. + +The benchmark uses a zero delta so repeated in-place launches remain stable. + +Example: + +.. code-block:: console + + .venv/bin/python benchmarks/kernels/benchmark_kimi_k3_attn_res.py \ + --output /tmp/attn_res.json +""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Callable +from dataclasses import asdict, dataclass +from functools import partial +from pathlib import Path + +import torch + +import vllm._custom_ops # noqa: F401 - loads torch.ops._C +from vllm.models.kimi_k3.nvidia.ops.attn_res import _attn_res_kernel +from vllm.platforms import current_platform +from vllm.triton_utils import triton + +HIDDEN_SIZE = 7168 +MAX_BLOCKS = 8 +EPS = 1e-5 + + +@dataclass +class Inputs: + prefix: torch.Tensor + delta: torch.Tensor + blocks: torch.Tensor + norm_weight: torch.Tensor + qk_weight: torch.Tensor + output_norm_weight: torch.Tensor + + +@dataclass +class Result: + num_tokens: int + num_blocks: int + backend: str + median_us: float + p20_us: float + p80_us: float + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--num-tokens", + type=int, + nargs="+", + default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192], + ) + parser.add_argument("--num-blocks", type=int, nargs="+", default=[1, 4, 8]) + parser.add_argument("--warmup-ms", type=int, default=50) + parser.add_argument("--repeat-ms", type=int, default=200) + parser.add_argument("--output", type=Path) + return parser.parse_args() + + +def make_inputs(num_tokens: int) -> Inputs: + prefix = torch.randn(num_tokens, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + blocks = torch.randn( + num_tokens, + MAX_BLOCKS, + HIDDEN_SIZE, + device="cuda", + dtype=torch.bfloat16, + ) + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = ( + torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5 + ) + output_norm_weight = 1 + 0.1 * torch.randn( + HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + return Inputs( + prefix=prefix, + delta=torch.zeros_like(prefix), + blocks=blocks, + norm_weight=norm_weight, + qk_weight=qk_weight, + output_norm_weight=output_norm_weight, + ) + + +def launch_triton(inputs: Inputs, output: torch.Tensor, num_blocks: int) -> None: + num_tokens = inputs.prefix.shape[0] + if num_tokens >= 256 or num_blocks <= 1: + block_l, num_warps = 1, 4 + else: + block_l, num_warps = 4, 8 + _attn_res_kernel[(num_tokens,)]( + inputs.prefix, + inputs.delta, + inputs.blocks, + inputs.norm_weight, + inputs.qk_weight, + inputs.output_norm_weight, + output, + inputs.prefix.stride(0), + inputs.delta.stride(0), + inputs.blocks.stride(0), + inputs.blocks.stride(1), + output.stride(0), + num_blocks, + HIDDEN_SIZE, + -1, + EPS, + EPS, + HAS_DELTA=True, + WRITE_BLOCK=False, + APPLY_OUTPUT_NORM=True, + BLOCK_L=block_l, + BLOCK_D=triton.next_power_of_2(HIDDEN_SIZE), + num_warps=num_warps, + num_stages=2, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + + +def launch_cuda(inputs: Inputs, output: torch.Tensor, num_blocks: int) -> None: + torch.ops._C.kimi_k3_attn_res( + inputs.prefix, + inputs.delta, + inputs.blocks, + inputs.norm_weight, + inputs.qk_weight, + inputs.output_norm_weight, + output, + num_blocks, + -1, + EPS, + EPS, + ) + + +def benchmark( + function: Callable[[], None], warmup_ms: int, repeat_ms: int +) -> tuple[float, float, float]: + median_ms, p20_ms, p80_ms = triton.testing.do_bench( + function, + warmup=warmup_ms, + rep=repeat_ms, + quantiles=[0.5, 0.2, 0.8], + ) + return median_ms * 1000, p20_ms * 1000, p80_ms * 1000 + + +def validate(inputs: Inputs, num_blocks: int) -> None: + outputs = [torch.empty_like(inputs.prefix) for _ in range(2)] + launch_triton(inputs, outputs[0], num_blocks) + launch_cuda(inputs, outputs[1], num_blocks) + torch.accelerator.synchronize() + torch.testing.assert_close(outputs[1], outputs[0], atol=8e-2, rtol=3e-2) + + +def main() -> None: + args = parse_args() + if torch.cuda.get_device_capability()[0] != 10: + raise SystemExit("The vLLM CUDA AttnRes kernel requires the SM100 family") + if not hasattr(torch.ops._C, "kimi_k3_attn_res"): + raise SystemExit("This vLLM build does not contain kimi_k3_attn_res") + if any(not 0 <= blocks <= MAX_BLOCKS for blocks in args.num_blocks): + raise SystemExit(f"--num-blocks values must be in [0, {MAX_BLOCKS}]") + + torch.manual_seed(0) + check_inputs = make_inputs(5) + for num_blocks in args.num_blocks: + validate(check_inputs, num_blocks) + print("correctness: all selected block counts passed", flush=True) + + results: list[Result] = [] + for num_blocks in args.num_blocks: + for num_tokens in args.num_tokens: + inputs = make_inputs(num_tokens) + output = torch.empty_like(inputs.prefix) + candidates = { + "vllm_triton": partial(launch_triton, inputs, output, num_blocks), + "vllm_cuda": partial(launch_cuda, inputs, output, num_blocks), + } + row: dict[str, float] = {} + for backend, candidate in candidates.items(): + median_us, p20_us, p80_us = benchmark( + candidate, args.warmup_ms, args.repeat_ms + ) + results.append( + Result( + num_tokens=num_tokens, + num_blocks=num_blocks, + backend=backend, + median_us=median_us, + p20_us=p20_us, + p80_us=p80_us, + ) + ) + row[backend] = median_us + print( + f"blocks={num_blocks} tokens={num_tokens:5d} " + f"triton={row['vllm_triton']:9.2f} us " + f"cuda={row['vllm_cuda']:9.2f} us", + flush=True, + ) + + if args.output is not None: + metadata = { + "device": torch.cuda.get_device_name(), + "capability": torch.cuda.get_device_capability(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "warmup_ms": args.warmup_ms, + "repeat_ms": args.repeat_ms, + } + args.output.write_text( + json.dumps( + {"metadata": metadata, "results": [asdict(item) for item in results]}, + indent=2, + ) + + "\n" + ) + print(f"wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu index 8c8f81380b..9eb040d6e5 100644 --- a/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu +++ b/csrc/libtorch_stable/kimi_k3/attn_res_kernel.cu @@ -10,7 +10,7 @@ // - Q=res_weight*rms_weight remains in registers across persistent tokens. // - V rows are converted once and cached as FP32 in TMEM between passes. // -// Integration contract: Kimi K3 H=7168, 1<=num_blocks<=8, and token-major +// Integration contract: Kimi K3 H=7168, 0<=num_blocks<=8, and token-major // block residual storage. #include "../torch_utils.h" @@ -230,14 +230,15 @@ __device__ __forceinline__ void cp_async_bulk(void* smem_dst, } template + bool RELEASE_TMEM = false, bool STATIC_COMMON_PATH = false, + bool HAS_OUTPUT_NORM = false> __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( - const bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res, + bf16_t* __restrict__ block_res, bf16_t* __restrict__ layer_res, const bf16_t* __restrict__ delta, const bf16_t* __restrict__ res_w, const bf16_t* __restrict__ rms_w, bf16_t* __restrict__ output, int T, int block_stride_m, int block_stride_r, float rms_eps, - const bf16_t* __restrict__ output_norm_weight, float output_norm_eps) { + const bf16_t* __restrict__ output_norm_weight, float output_norm_eps, + int block_write_idx) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100 constexpr float LOG2_E = 1.4426950408889634f; constexpr int N_CHUNK = NC; @@ -254,12 +255,15 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( static_assert(TMEM_V_COLS_TOTAL <= TMEM_COLS_ALLOC); static_assert(H >= 4096 && H <= 8192); static_assert(H % K_TILE == 0); + static_assert(!STATIC_COMMON_PATH || HAS_OUTPUT_NORM); const int tid = threadIdx.x; const int wid = tid >> 5; const int lane = tid & 31; const int TB = T * B; const int num_ctas = gridDim.x; + const bool has_delta = STATIC_COMMON_PATH || delta != nullptr; + const int write_idx = STATIC_COMMON_PATH ? -1 : block_write_idx; constexpr int num_chunks = (N + N_CHUNK - 1) / N_CHUNK; const int comp_wid = wid - 1; @@ -270,10 +274,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( const int k_local = ct_in_group * VEC; constexpr size_t V_BYTES = (size_t)NUM_BUFS * H * sizeof(bf16_t); - constexpr size_t DELTA_BYTES = - HAS_DELTA ? (size_t)CHUNK_DEPTH * H * sizeof(bf16_t) : 0; + constexpr size_t DELTA_BYTES = (size_t)CHUNK_DEPTH * H * sizeof(bf16_t); constexpr size_t OUTPUT_NORM_BYTES = - OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0; + HAS_OUTPUT_NORM ? (size_t)H * sizeof(bf16_t) : 0; extern __shared__ __align__(16) char smem_raw[]; bf16_t* v_bufs = reinterpret_cast(smem_raw); // [NUM_BUFS][H] bf16_t* delta_bufs = reinterpret_cast(smem_raw + V_BYTES); @@ -297,15 +300,12 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( mbarrier_init(plan.bar_ready[i], 1); mbarrier_init(plan.bar_consumed[i], CONSUMER_WARPS); } - if constexpr (OUTPUT_NORM_IN_SMEM) { + if constexpr (HAS_OUTPUT_NORM) { mbarrier_init(plan.bar_output_norm_ready, 1); } fence_mbarrier_init(); } - // gdc wait BEFORE tmem alloc - cudaGridDependencySynchronize(); - if (wid == 1) { tmem_allocate(TMEM_COLS_ALLOC, &plan.tmem_base); if constexpr (RELEASE_TMEM) { @@ -314,7 +314,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( } __syncthreads(); - if constexpr (OUTPUT_NORM_IN_SMEM) { + if constexpr (HAS_OUTPUT_NORM) { if (wid == 0 && elect_one_sync()) { mbarrier_expect_tx(plan.bar_output_norm_ready, H * (int)sizeof(bf16_t)); cp_async_bulk(output_norm_buf, output_norm_weight, H * sizeof(bf16_t), @@ -365,6 +365,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( if (wid == 0) { if (elect_one_sync()) { long long gci = 0; + bool pdl_waited = false; for (int tb = blockIdx.x; tb < TB; tb += num_ctas) { const int t = tb / B; for (int ci = 0; ci < num_chunks; ci++, gci++) { @@ -374,7 +375,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( int pc = phase_of(gci); mbarrier_wait(plan.bar_consumed[chunk_slot], pc ^ 1); int transaction_bytes = an * H * (int)sizeof(bf16_t); - if constexpr (HAS_DELTA) { + if (has_delta) { int prefix_n = N - 1 - ns; if (prefix_n >= 0 && prefix_n < an) { transaction_bytes += H * (int)sizeof(bf16_t); @@ -388,12 +389,20 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( const bf16_t* src = residual_addr(block_res, layer_res, ns + n, N, t, block_stride_m, block_stride_r, H); + if (!has_delta && ns + n == N - 1 && !pdl_waited) { + cudaGridDependencySynchronize(); + pdl_waited = true; + } cp_async_bulk(buf_ptr(slot), src, H * sizeof(bf16_t), plan.bar_ready[chunk_slot]); } - if constexpr (HAS_DELTA) { + if (has_delta) { int prefix_n = N - 1 - ns; if (prefix_n >= 0 && prefix_n < an) { + if (!pdl_waited) { + cudaGridDependencySynchronize(); + pdl_waited = true; + } cp_async_bulk(delta_buf_ptr(chunk_slot), delta + (long long)tb * H, H * sizeof(bf16_t), plan.bar_ready[chunk_slot]); @@ -401,6 +410,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( } } } + cudaTriggerProgrammaticLaunchCompletion(); } } else { float acc32[ACC_PER_THREAD] = {}; @@ -442,9 +452,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( int2 vp = *reinterpret_cast(buf_ptr(slot) + h_base); auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); - if constexpr (HAS_DELTA) { - int prefix_n = N - 1 - ns; - if (n == prefix_n) { + int prefix_n = N - 1 - ns; + if (n == prefix_n && (has_delta || write_idx >= 0)) { + if (has_delta) { const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base; #pragma unroll @@ -456,6 +466,11 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( *reinterpret_cast(layer_res + (long long)tb * H + h_base) = vp; } + if (write_idx >= 0) { + *reinterpret_cast( + block_res + (long long)(tb / B) * block_stride_m + + write_idx * block_stride_r + h_base) = vp; + } } float2 f[2] = {__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1])}; @@ -480,9 +495,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( int slot = slot_of(gci, n); int4 vp = *reinterpret_cast(buf_ptr(slot) + h_base); auto* v2 = reinterpret_cast<__nv_bfloat162*>(&vp); - if constexpr (HAS_DELTA) { - int prefix_n = N - 1 - ns; - if (n == prefix_n) { + int prefix_n = N - 1 - ns; + if (n == prefix_n && (has_delta || write_idx >= 0)) { + if (has_delta) { const bf16_t* delta_ptr = delta_buf_ptr(chunk_slot) + h_base; #pragma unroll for (int j = 0; j < VEC / 2; j++) { @@ -493,6 +508,11 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( *reinterpret_cast(layer_res + (long long)tb * H + h_base) = vp; } + if (write_idx >= 0) { + *reinterpret_cast( + block_res + (long long)(tb / B) * block_stride_m + + write_idx * block_stride_r + h_base) = vp; + } } float2 f[4] = { __bfloat1622float2(v2[0]), __bfloat1622float2(v2[1]), @@ -791,11 +811,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( } if constexpr (HAS_OUTPUT_NORM) { - if constexpr (OUTPUT_NORM_IN_SMEM) { - // The immutable weight copy is acquired once, at its first use. - if (tb == blockIdx.x) { - mbarrier_wait(plan.bar_output_norm_ready, 0); - } + // The immutable weight copy is acquired once, at its first use. + if (tb == blockIdx.x) { + mbarrier_wait(plan.bar_output_norm_ready, 0); } float output_sq = output_sq_pair.x + output_sq_pair.y; #pragma unroll @@ -826,9 +844,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( auto* values = reinterpret_cast(&packed); #pragma unroll for (int j = 0; j < 4; j++) { - const bf16_t* weight_ptr = - OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; - float weight = __bfloat162float(weight_ptr[h_base + j]); + float weight = __bfloat162float(output_norm_buf[h_base + j]); values[j] = __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight); } @@ -843,9 +859,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( auto* values = reinterpret_cast(&packed); #pragma unroll for (int j = 0; j < VEC; j++) { - const bf16_t* weight_ptr = - OUTPUT_NORM_IN_SMEM ? output_norm_buf : output_norm_weight; - float weight = __bfloat162float(weight_ptr[h_base + j]); + float weight = __bfloat162float(output_norm_buf[h_base + j]); values[j] = __float2bfloat16(acc32[si * VEC + j] * output_rsigma * weight); } @@ -855,7 +869,6 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( } } - cudaTriggerProgrammaticLaunchCompletion(); __syncthreads(); if (wid == 1) { tmem_free(plan.tmem_base, TMEM_COLS_ALLOC); @@ -868,23 +881,22 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel( } template -static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual, + bool STATIC_COMMON_PATH = false, bool HAS_OUTPUT_NORM = false> +static void launch_fwd(bf16_t* block_residual, bf16_t* layer_residual, const bf16_t* delta, const bf16_t* res_weight, const bf16_t* rms_weight, bf16_t* output, int T, float rms_eps, int num_sm, cudaStream_t stream, const bf16_t* output_norm_weight = nullptr, float output_norm_eps = 0.f, int block_stride_m = 0, - int block_stride_r = 0) { + int block_stride_r = 0, int block_write_idx = -1) { constexpr size_t smem_size = - ((size_t)CHUNK_DEPTH * (NC + (HAS_DELTA ? 1 : 0)) * H * sizeof(bf16_t) + - (OUTPUT_NORM_IN_SMEM ? (size_t)H * sizeof(bf16_t) : 0) + + ((size_t)(CHUNK_DEPTH * (NC + 1) + (HAS_OUTPUT_NORM ? 1 : 0)) * H * + sizeof(bf16_t) + sizeof(FwdSmemPlan) + 15) & ~size_t(15); auto kernel = - &attn_res_fwd_online_v2_kernel; + &attn_res_fwd_online_v2_kernel; static bool attrs_set = false; if (!attrs_set) { if (smem_size > 48 * 1024) { @@ -907,20 +919,21 @@ static void launch_fwd(const bf16_t* block_residual, bf16_t* layer_residual, cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, delta, res_weight, rms_weight, output, T, block_stride_m, block_stride_r, rms_eps, output_norm_weight, - output_norm_eps); + output_norm_eps, block_write_idx); } } // namespace fwd_prod_v2 } // namespace sm100 void kimi_k3_attn_res(torch::stable::Tensor& prefix, - torch::stable::Tensor const& delta, - torch::stable::Tensor const& blocks, + std::optional delta, + torch::stable::Tensor& blocks, torch::stable::Tensor const& norm_weight, torch::stable::Tensor const& qk_weight, - torch::stable::Tensor const& output_norm_weight, + std::optional output_norm_weight, torch::stable::Tensor& output, int64_t num_blocks, - double eps, double output_norm_eps) { + int64_t block_write_idx, double eps, + double output_norm_eps) { int const num_tokens = static_cast(prefix.size(0)); int const device = prefix.get_device_index(); torch::stable::accelerator::DeviceGuard const device_guard(device); @@ -928,49 +941,57 @@ void kimi_k3_attn_res(torch::stable::Tensor& prefix, STD_TORCH_CHECK(properties->major == 10, "Kimi K3 AttnRes requires the SM100 family"); int64_t const hidden_size = prefix.size(1); + STD_TORCH_CHECK(hidden_size == 7168, + "Kimi K3 AttnRes requires hidden_size=7168, got ", + hidden_size); STD_TORCH_CHECK(prefix.stride(0) == hidden_size && - delta.stride(0) == hidden_size && + (!delta.has_value() || delta->stride(0) == hidden_size) && output.stride(0) == hidden_size, "Kimi K3 AttnRes requires densely packed rows; got strides ", - prefix.stride(0), ", ", delta.stride(0), ", ", + prefix.stride(0), ", ", + delta.has_value() ? delta->stride(0) : 0, ", ", output.stride(0), " for hidden_size ", hidden_size); + STD_TORCH_CHECK( + num_blocks >= 0 && num_blocks <= 8 && num_blocks <= blocks.size(1), + "Kimi K3 AttnRes: num_blocks must be 0..min(8, blocks.size(1))"); + STD_TORCH_CHECK(block_write_idx >= -1 && block_write_idx < blocks.size(1), + "Kimi K3 AttnRes: block_write_idx must be -1 or a valid " + "block row, got ", + block_write_idx); + + if (num_tokens == 0) { + return; + } using namespace sm100::fwd_prod_v2; - // Two-source chunks and two resident CTAs are beneficial once setup is - // amortized by the long, full eight-block prefill workload. - if (num_blocks == 8 && num_tokens >= 4096) { - launch_fwd<7168, 9, 2, true, true, true, true>( - static_cast(blocks.data_ptr()), - static_cast(prefix.data_ptr()), - static_cast(delta.data_ptr()), - static_cast(qk_weight.data_ptr()), - static_cast(norm_weight.data_ptr()), - static_cast(output.data_ptr()), num_tokens, - static_cast(eps), properties->multiProcessorCount, - get_current_cuda_stream(device), - static_cast(output_norm_weight.data_ptr()), - static_cast(output_norm_eps), static_cast(blocks.stride(0)), - static_cast(blocks.stride(1))); - } else { - auto dispatch = [&](auto nsrc_tok) { - constexpr int NSRC = decltype(nsrc_tok)::value; - launch_fwd<7168, NSRC, 3, false, true, true, true>( - static_cast(blocks.data_ptr()), + cudaStream_t stream = get_current_cuda_stream(device); + auto dispatch_blocks = [&](auto static_common_path, auto has_output_norm) { + constexpr bool STATIC_COMMON_PATH = decltype(static_common_path)::value; + constexpr bool HAS_OUTPUT_NORM = decltype(has_output_norm)::value; + auto dispatch = [&](auto source_count) { + constexpr int NUM_SOURCES = decltype(source_count)::value; + launch_fwd<7168, NUM_SOURCES, 3, false, STATIC_COMMON_PATH, + HAS_OUTPUT_NORM>( + static_cast(blocks.data_ptr()), static_cast(prefix.data_ptr()), - static_cast(delta.data_ptr()), + delta.has_value() ? static_cast(delta->data_ptr()) + : nullptr, static_cast(qk_weight.data_ptr()), static_cast(norm_weight.data_ptr()), static_cast(output.data_ptr()), num_tokens, - static_cast(eps), properties->multiProcessorCount, - get_current_cuda_stream(device), - static_cast(output_norm_weight.data_ptr()), + static_cast(eps), properties->multiProcessorCount, stream, + output_norm_weight.has_value() + ? static_cast(output_norm_weight->data_ptr()) + : nullptr, static_cast(output_norm_eps), static_cast(blocks.stride(0)), - static_cast(blocks.stride(1))); + static_cast(blocks.stride(1)), + static_cast(block_write_idx)); }; - // The source count is num_blocks + 1; specialising on it makes the chunk - // count, per-chunk source count and prefix-chunk test compile-time. switch (num_blocks) { + case 0: + dispatch(std::integral_constant{}); + break; case 1: dispatch(std::integral_constant{}); break; @@ -995,9 +1016,15 @@ void kimi_k3_attn_res(torch::stable::Tensor& prefix, case 8: dispatch(std::integral_constant{}); break; - default: - STD_TORCH_CHECK(false, "Kimi K3 AttnRes: num_blocks must be 1..8"); } + }; + if (delta.has_value() && output_norm_weight.has_value() && + block_write_idx < 0) { + dispatch_blocks(std::true_type{}, std::true_type{}); + } else if (output_norm_weight.has_value()) { + dispatch_blocks(std::false_type{}, std::true_type{}); + } else { + dispatch_blocks(std::false_type{}, std::false_type{}); } cudaError_t const error = cudaGetLastError(); STD_TORCH_CHECK( diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 3b79fc591f..08e5fe3226 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -411,13 +411,14 @@ void fused_gdn_decode_post_conv_mtp( #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES void kimi_k3_attn_res(torch::stable::Tensor& prefix, - torch::stable::Tensor const& delta, - torch::stable::Tensor const& blocks, + std::optional delta, + torch::stable::Tensor& blocks, torch::stable::Tensor const& norm_weight, torch::stable::Tensor const& qk_weight, - torch::stable::Tensor const& output_norm_weight, + std::optional output_norm_weight, torch::stable::Tensor& output, int64_t num_blocks, - double eps, double output_norm_eps); + int64_t block_write_idx, double eps, + double output_norm_eps); #endif // Sampler kernels (shared CUDA/ROCm) diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 85af95e67e..2247281e72 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -542,9 +542,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES ops.def( "kimi_k3_attn_res(" - "Tensor! prefix, Tensor delta, Tensor blocks, Tensor norm_weight, " - "Tensor qk_weight, Tensor output_norm_weight, Tensor! output, " - "int num_blocks, float eps, float output_norm_eps) -> ()"); + "Tensor! prefix, Tensor? delta, Tensor! blocks, Tensor norm_weight, " + "Tensor qk_weight, Tensor? output_norm_weight, Tensor! output, " + "int num_blocks, " + "int block_write_idx, float eps, " + "float output_norm_eps) -> ()"); #endif // Apply repetition penalties to logits in-place. diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py index 38a34bc923..574f7a5925 100644 --- a/tests/models/kimi_k3/test_attn_res.py +++ b/tests/models/kimi_k3/test_attn_res.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib + import pytest import torch import torch.nn.functional as F @@ -12,6 +14,7 @@ from vllm.platforms import current_platform HIDDEN_SIZE = 7168 MAX_BLOCKS = 8 EPS = 1e-5 +attn_res_module = importlib.import_module("vllm.models.kimi_k3.nvidia.ops.attn_res") def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor: @@ -60,6 +63,7 @@ def _reference( pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"), pytest.param(3, 8, 0, False, False, "triton", id="triton-full"), pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"), + pytest.param(1, 4, 0, False, True, "nvidia", id="nvidia-single-token"), pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"), pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"), pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"), @@ -194,6 +198,55 @@ def test_attn_res_without_output_norm(): torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2) +@pytest.mark.parametrize( + ("num_blocks", "has_delta", "has_output_norm", "write_block"), + [ + (0, False, True, True), + (4, True, True, True), + (4, False, True, False), + (8, True, False, False), + ], +) +def test_sm100_variants_do_not_fall_back_to_triton( + monkeypatch, + num_blocks: int, + has_delta: bool, + has_output_norm: bool, + write_block: bool, +): + if not current_platform.is_device_capability_family(100) or not hasattr( + torch.ops._C, "kimi_k3_attn_res" + ): + pytest.skip("native AttnRes requires an SM100 build") + + class FailingKernel: + def __getitem__(self, grid): + raise AssertionError(f"unexpected Triton launch with grid {grid}") + + monkeypatch.setattr(attn_res_module, "_attn_res_kernel", FailingKernel()) + prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + delta = torch.zeros_like(prefix) if has_delta else None + blocks = torch.randn( + 1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16 + ) + norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + qk_weight = torch.randn_like(norm_weight) / HIDDEN_SIZE**0.5 + output_norm_weight = torch.ones_like(norm_weight) if has_output_norm else None + + attn_res( + prefix, + delta, + blocks, + norm_weight, + qk_weight, + output_norm_weight, + num_blocks, + num_blocks if write_block else -1, + EPS, + EPS, + ) + + @pytest.mark.parametrize("num_tokens", [0, 1, 17]) def test_fused_mtp_input(num_tokens: int): positions = torch.arange(num_tokens, device="cuda") diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 78a4c43159..080893d648 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2904,12 +2904,13 @@ def concat_and_cache_mla_grouped( def kimi_k3_attn_res( prefix: torch.Tensor, - delta: torch.Tensor, + delta: torch.Tensor | None, blocks: torch.Tensor, norm_weight: torch.Tensor, qk_weight: torch.Tensor, - output_norm_weight: torch.Tensor, + output_norm_weight: torch.Tensor | None, num_blocks: int, + block_write_idx: int, eps: float, output_norm_eps: float, ) -> torch.Tensor: @@ -2923,6 +2924,7 @@ def kimi_k3_attn_res( output_norm_weight, output, num_blocks, + block_write_idx, eps, output_norm_eps, ) diff --git a/vllm/models/kimi_k3/nvidia/ops/attn_res.py b/vllm/models/kimi_k3/nvidia/ops/attn_res.py index 06f73ce8f7..06a993c0a1 100644 --- a/vllm/models/kimi_k3/nvidia/ops/attn_res.py +++ b/vllm/models/kimi_k3/nvidia/ops/attn_res.py @@ -186,16 +186,16 @@ def attn_res( assert norm_weight.stride(-1) == 1 assert qk_weight.stride(-1) == 1 assert output_norm_weight is None or output_norm_weight.stride(-1) == 1 - # The in-tree NVIDIA kernel covers the common fused-add + output-norm path; - # Triton handles block boundaries and final pre-norm output. The native op - # is only compiled for SM100 under CUDA >= 13, so a device check alone is - # not enough to know it exists. + # The native kernel covers every Kimi-K3 AttnRes variant on dense SM100 + # inputs. The op is only compiled under CUDA >= 13, so a device check alone + # is not enough to know it exists. if ( hidden_size == 7168 - and delta is not None - and output_norm_weight is not None - and num_blocks > 0 - and block_write_idx < 0 + and prefix.stride(0) == hidden_size + and (delta is None or delta.stride(0) == hidden_size) + and 0 <= num_blocks <= 8 + and num_blocks <= blocks.shape[1] + and -1 <= block_write_idx < blocks.shape[1] and current_platform.is_device_capability_family(100) and hasattr(torch.ops._C, "kimi_k3_attn_res") ): @@ -207,6 +207,7 @@ def attn_res( qk_weight, output_norm_weight, num_blocks, + block_write_idx, eps, output_norm_eps, )