From d75cab85629e34d3a905da3ea205e6bfae5751b5 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Fri, 5 Jun 2026 17:57:01 +0800 Subject: [PATCH 01/35] Add Qwen3.5 decode kernel scaffolding --- csrc/qwen35/decode/qwen35_conv1d_decode.cu | 29 ++ csrc/qwen35/decode/qwen35_decode_common.cuh | 73 +++ csrc/qwen35/decode/qwen35_layout_decode.cu | 132 +++++ csrc/qwen35/decode/qwen35_layout_kernel.hpp | 127 +++++ .../qwen35/decode/qwen35_scalar_kda_decode.cu | 134 ++++++ .../decode/qwen35_scalar_kda_kernel.hpp | 247 ++++++++++ .../decode/qwen35_scalar_kda_mainloop.hpp | 452 ++++++++++++++++++ cula/ops/qwen35_conv1d_decode.py | 170 +++++++ cula/ops/qwen35_conv1d_prefill.py | 36 ++ cula/ops/qwen35_scalar_kda_decode.py | 75 +++ cula/ops/qwen35_scalar_kda_prefill.py | 37 ++ cula/qwen35/__init__.py | 27 ++ cula/qwen35/common.py | 135 ++++++ cula/qwen35/runtime.py | 165 +++++++ docs/qwen35_kernel_plan.md | 40 ++ 15 files changed, 1879 insertions(+) create mode 100644 csrc/qwen35/decode/qwen35_conv1d_decode.cu create mode 100644 csrc/qwen35/decode/qwen35_decode_common.cuh create mode 100644 csrc/qwen35/decode/qwen35_layout_decode.cu create mode 100644 csrc/qwen35/decode/qwen35_layout_kernel.hpp create mode 100644 csrc/qwen35/decode/qwen35_scalar_kda_decode.cu create mode 100644 csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp create mode 100644 csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp create mode 100644 cula/ops/qwen35_conv1d_decode.py create mode 100644 cula/ops/qwen35_conv1d_prefill.py create mode 100644 cula/ops/qwen35_scalar_kda_decode.py create mode 100644 cula/ops/qwen35_scalar_kda_prefill.py create mode 100644 cula/qwen35/__init__.py create mode 100644 cula/qwen35/common.py create mode 100644 cula/qwen35/runtime.py create mode 100644 docs/qwen35_kernel_plan.md diff --git a/csrc/qwen35/decode/qwen35_conv1d_decode.cu b/csrc/qwen35/decode/qwen35_conv1d_decode.cu new file mode 100644 index 00000000..cbce4896 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_conv1d_decode.cu @@ -0,0 +1,29 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qwen35_decode_common.cuh" + +#include + +namespace cula::qwen35::decode { + +void run_qwen35_conv1d_decode(ConvDecodeParams& params) { + (void)params; + TORCH_CHECK( + false, + "run_qwen35_conv1d_decode is not implemented yet. " + "Planned kernel: single-token depthwise causal conv1d + silu for 10240 channels."); +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_decode_common.cuh b/csrc/qwen35/decode/qwen35_decode_common.cuh new file mode 100644 index 00000000..53b64d65 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_decode_common.cuh @@ -0,0 +1,73 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include + +namespace cula::qwen35::decode { + +inline constexpr int kNumQKHeads = 16; +inline constexpr int kNumVHeads = 48; +inline constexpr int kHeadDimQK = 128; +inline constexpr int kHeadDimV = 128; +inline constexpr int kConvKernelSize = 4; +inline constexpr int kQDim = kNumQKHeads * kHeadDimQK; +inline constexpr int kKDim = kNumQKHeads * kHeadDimQK; +inline constexpr int kVDim = kNumVHeads * kHeadDimV; +inline constexpr int kMixedQKVDim = kQDim + kKDim + kVDim; + +struct ConvDecodeParams { + at::Tensor mixed_qkv; // [B, 1, 10240] + at::Tensor conv_state; // [B, 10240, 4] + at::Tensor conv_weight; // [10240, 4] + at::Tensor out; // [B, 1, 10240] +}; + +struct LayoutDecodeParams { + at::Tensor mixed_qkv_conv; // [N, 10240] + at::Tensor a; // [N, 48] + at::Tensor b; // [N, 48] + at::Tensor q_rep; // [N, 48, 128] + at::Tensor k_rep; // [N, 48, 128] + at::Tensor v; // [N, 48, 128] + at::Tensor a_kernel; // [N, 48] + at::Tensor b_kernel; // [N, 48] +}; + +struct ScalarKdaDecodeParams { + // Dtype contract for the first implementation: + // - activations / outputs: half or bf16 + // q_rep, k_rep, v, a_kernel, b_kernel, out + // - recurrent parameters / state: float32 + // A_log, dt_bias, recurrent_state + at::Tensor q_rep; // [N, 48, 128] + at::Tensor k_rep; // [N, 48, 128] + at::Tensor v; // [N, 48, 128] + at::Tensor a_kernel; // [N, 48] + at::Tensor b_kernel; // [N, 48] + at::Tensor A_log; // [48], float32 + at::Tensor dt_bias; // [48], float32 + at::Tensor recurrent_state; // [pool, 48, 128, 128], float32 + at::Tensor pool_idx; // [N], int32 + at::Tensor out; // [N, 48, 128] +}; + +void run_qwen35_conv1d_decode(ConvDecodeParams& params); +void run_qwen35_layout_decode(LayoutDecodeParams& params); +void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params); + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_layout_decode.cu b/csrc/qwen35/decode/qwen35_layout_decode.cu new file mode 100644 index 00000000..e601f17f --- /dev/null +++ b/csrc/qwen35/decode/qwen35_layout_decode.cu @@ -0,0 +1,132 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qwen35_decode_common.cuh" +#include "qwen35_layout_kernel.hpp" + +#include +#include +#include +#include +#include +#include + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +void check_tensor_shape_2d(const at::Tensor& tensor, const char* name) { + TORCH_CHECK( + tensor.dim() == 2, + name, + " must have rank 2, but got rank ", + tensor.dim(), + "."); +} + +} // namespace + +namespace cula::qwen35::decode { + +void run_qwen35_layout_decode(LayoutDecodeParams& params) { + const at::Tensor& mixed_qkv_conv = params.mixed_qkv_conv; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& q_rep = params.q_rep; + const at::Tensor& k_rep = params.k_rep; + const at::Tensor& v = params.v; + const at::Tensor& a_kernel = params.a_kernel; + const at::Tensor& b_kernel = params.b_kernel; + + TORCH_CHECK(mixed_qkv_conv.is_cuda(), "mixed_qkv_conv must be a CUDA tensor."); + TORCH_CHECK(mixed_qkv_conv.is_contiguous(), "mixed_qkv_conv must be contiguous."); + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == a.scalar_type() && + mixed_qkv_conv.scalar_type() == b.scalar_type() && + mixed_qkv_conv.scalar_type() == q_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == k_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == v.scalar_type() && + mixed_qkv_conv.scalar_type() == a_kernel.scalar_type() && + mixed_qkv_conv.scalar_type() == b_kernel.scalar_type(), + "All layout decode tensors must share the same dtype."); + + check_tensor_shape_2d(a, "a"); + check_tensor_shape_2d(b, "b"); + + TORCH_CHECK( + mixed_qkv_conv.dim() == 2 && mixed_qkv_conv.size(1) == kMixedQKVDim, + "mixed_qkv_conv must have shape [N, 10240]."); + + const int64_t batch_size = mixed_qkv_conv.size(0); + const at::Device device = mixed_qkv_conv.device(); + + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(q_rep, "q_rep", device); + check_tensor_device(k_rep, "k_rep", device); + check_tensor_device(v, "v", device); + check_tensor_device(a_kernel, "a_kernel", device); + check_tensor_device(b_kernel, "b_kernel", device); + + TORCH_CHECK(q_rep.is_contiguous(), "q_rep must be contiguous."); + TORCH_CHECK(k_rep.is_contiguous(), "k_rep must be contiguous."); + TORCH_CHECK(v.is_contiguous(), "v must be contiguous."); + TORCH_CHECK(a_kernel.is_contiguous(), "a_kernel must be contiguous."); + TORCH_CHECK(b_kernel.is_contiguous(), "b_kernel must be contiguous."); + + TORCH_CHECK( + q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({batch_size, kNumVHeads, kHeadDimQK}), + "q_rep must have shape [N, 48, 128]."); + TORCH_CHECK( + k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({batch_size, kNumVHeads, kHeadDimQK}), + "k_rep must have shape [N, 48, 128]."); + TORCH_CHECK( + v.dim() == 3 && v.sizes() == at::IntArrayRef({batch_size, kNumVHeads, kHeadDimV}), + "v must have shape [N, 48, 128]."); + TORCH_CHECK( + a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), + "a_kernel must have shape [N, 48]."); + TORCH_CHECK( + b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), + "b_kernel must have shape [N, 48]."); + + TORCH_CHECK(a.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), "a must have shape [N, 48]."); + TORCH_CHECK(b.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), "b must have shape [N, 48]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + constexpr int threads = 32; + dim3 grid(kNumVHeads, static_cast(batch_size), 1); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + mixed_qkv_conv.scalar_type(), + "qwen35_layout_decode_kernel_cute", + [&] { + qwen35_layout_decode_kernel_cute<<>>( + mixed_qkv_conv.data_ptr(), + a.data_ptr(), + b.data_ptr(), + q_rep.data_ptr(), + k_rep.data_ptr(), + v.data_ptr(), + a_kernel.data_ptr(), + b_kernel.data_ptr(), + batch_size); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_layout_kernel.hpp b/csrc/qwen35/decode/qwen35_layout_kernel.hpp new file mode 100644 index 00000000..36e9373e --- /dev/null +++ b/csrc/qwen35/decode/qwen35_layout_kernel.hpp @@ -0,0 +1,127 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "qwen35_decode_common.cuh" + +#include +#include +#include + +namespace cula::qwen35::decode { + +using namespace cute; + +template +CUTE_DEVICE void copy_vec_contiguous( + scalar_t* __restrict__ dst, + const scalar_t* __restrict__ src) { + constexpr int kBytes = sizeof(scalar_t) * kVec; + if constexpr (kBytes == 16 || kBytes == 8) { + using VecType = cutlass::AlignedArray; + auto dst_addr = reinterpret_cast(dst); + auto src_addr = reinterpret_cast(src); + if ((dst_addr % alignof(VecType) == 0) && (src_addr % alignof(VecType) == 0)) { + *reinterpret_cast(dst) = *reinterpret_cast(src); + return; + } + } + +#pragma unroll + for (int i = 0; i < kVec; ++i) { + dst[i] = src[i]; + } +} + +template +__global__ void qwen35_layout_decode_kernel_cute( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + scalar_t* __restrict__ q_rep, + scalar_t* __restrict__ k_rep, + scalar_t* __restrict__ v_out, + scalar_t* __restrict__ a_kernel, + scalar_t* __restrict__ b_kernel, + int64_t token_count) { + static_assert(kNumVHeads % kNumQKHeads == 0); + constexpr int kRepeatFactor = kNumVHeads / kNumQKHeads; + // TODO(qwen35-layout-opt): + // - Re-evaluate whether Vec=8 is profitable for bf16/fp16 on the target GPUs. + // - Push more of the q/k repeat mapping into compile-time CuTe layout transforms. + // - Revisit whether a shared-memory staging path is worthwhile after profiling. + // - Consider widening the a/b writeback path if it shows up in profiling. + constexpr int kVec = 4; + static_assert(kHeadDimV % kVec == 0); + static_assert(kHeadDimQK == kHeadDimV); + static_assert(kHeadDimQK % kVec == 0); + + const int token_idx = static_cast(blockIdx.y); + const int hv = static_cast(blockIdx.x); + const int tid = static_cast(threadIdx.x); + + if (token_idx >= token_count || hv >= kNumVHeads) { + return; + } + + const int mapped_h = hv / kRepeatFactor; + + auto qk_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto out_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + + const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kMixedQKVDim; + const scalar_t* q_src_ptr = token_ptr; + const scalar_t* k_src_ptr = token_ptr + kQDim; + const scalar_t* v_src_ptr = token_ptr + kQDim + kKDim; + + scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; + scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; + scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kNumVHeads * kHeadDimV; + + // Current version uses a direct GMEM->GMEM vector copy path. This keeps the + // kernel simple while already removing the scalar-copy bottleneck from the + // first draft. More aggressive staging/copy strategies should be driven by + // profiling rather than added pre-emptively. + for (int vec_idx = tid; vec_idx < kHeadDimV / kVec; vec_idx += blockDim.x) { + const int d = vec_idx * kVec; + const int q_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int k_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int v_src_idx = crd2idx(make_coord(hv, d), v_src_layout); + const int dst_idx = crd2idx(make_coord(hv, d), out_layout); + + copy_vec_contiguous(q_dst_ptr + dst_idx, q_src_ptr + q_src_idx); + copy_vec_contiguous(k_dst_ptr + dst_idx, k_src_ptr + k_src_idx); + copy_vec_contiguous(v_dst_ptr + dst_idx, v_src_ptr + v_src_idx); + } + + if (tid == 0) { + // TODO(qwen35-layout-opt): If a/b copy becomes measurable, fuse a wider + // per-head copy path here instead of scalar head writes. + const int head_idx = crd2idx(make_coord(hv), head_layout); + const int64_t token_head_offset = static_cast(token_idx) * kNumVHeads + head_idx; + a_kernel[token_head_offset] = a[token_head_offset]; + b_kernel[token_head_offset] = b[token_head_offset]; + } +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu new file mode 100644 index 00000000..abed9854 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu @@ -0,0 +1,134 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qwen35_decode_common.cuh" +#include "qwen35_scalar_kda_kernel.hpp" + +#include +#include +#include +#include + +namespace cula::qwen35::decode { + +namespace { + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +} // namespace + +void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params) { + const at::Tensor& q_rep = params.q_rep; + const at::Tensor& k_rep = params.k_rep; + const at::Tensor& v = params.v; + const at::Tensor& a_kernel = params.a_kernel; + const at::Tensor& b_kernel = params.b_kernel; + const at::Tensor& A_log = params.A_log; + const at::Tensor& dt_bias = params.dt_bias; + const at::Tensor& recurrent_state = params.recurrent_state; + const at::Tensor& pool_idx = params.pool_idx; + const at::Tensor& out = params.out; + + TORCH_CHECK(q_rep.is_cuda(), "q_rep must be a CUDA tensor."); + const at::Device device = q_rep.device(); + + check_tensor_device(k_rep, "k_rep", device); + check_tensor_device(v, "v", device); + check_tensor_device(a_kernel, "a_kernel", device); + check_tensor_device(b_kernel, "b_kernel", device); + check_tensor_device(A_log, "A_log", device); + check_tensor_device(dt_bias, "dt_bias", device); + check_tensor_device(recurrent_state, "recurrent_state", device); + check_tensor_device(pool_idx, "pool_idx", device); + check_tensor_device(out, "out", device); + + TORCH_CHECK(q_rep.is_contiguous(), "q_rep must be contiguous."); + TORCH_CHECK(k_rep.is_contiguous(), "k_rep must be contiguous."); + TORCH_CHECK(v.is_contiguous(), "v must be contiguous."); + TORCH_CHECK(a_kernel.is_contiguous(), "a_kernel must be contiguous."); + TORCH_CHECK(b_kernel.is_contiguous(), "b_kernel must be contiguous."); + TORCH_CHECK(A_log.is_contiguous(), "A_log must be contiguous."); + TORCH_CHECK(dt_bias.is_contiguous(), "dt_bias must be contiguous."); + TORCH_CHECK(recurrent_state.is_contiguous(), "recurrent_state must be contiguous."); + TORCH_CHECK(pool_idx.is_contiguous(), "pool_idx must be contiguous."); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous."); + + TORCH_CHECK( + q_rep.scalar_type() == k_rep.scalar_type() && q_rep.scalar_type() == v.scalar_type() && + q_rep.scalar_type() == a_kernel.scalar_type() && q_rep.scalar_type() == b_kernel.scalar_type() && + q_rep.scalar_type() == out.scalar_type(), + "q_rep/k_rep/v/a_kernel/b_kernel/out must share the same dtype."); + TORCH_CHECK(A_log.scalar_type() == at::kFloat, "A_log must be float32."); + TORCH_CHECK(dt_bias.scalar_type() == at::kFloat, "dt_bias must be float32."); + TORCH_CHECK(recurrent_state.scalar_type() == at::kFloat, "recurrent_state must be float32."); + TORCH_CHECK(pool_idx.scalar_type() == at::kInt, "pool_idx must be int32."); + + const int64_t token_count = q_rep.size(0); + TORCH_CHECK( + q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimQK}), + "q_rep must have shape [N, 48, 128]."); + TORCH_CHECK( + k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimQK}), + "k_rep must have shape [N, 48, 128]."); + TORCH_CHECK( + v.dim() == 3 && v.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimV}), + "v must have shape [N, 48, 128]."); + TORCH_CHECK( + a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({token_count, kNumVHeads}), + "a_kernel must have shape [N, 48]."); + TORCH_CHECK( + b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({token_count, kNumVHeads}), + "b_kernel must have shape [N, 48]."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == kNumVHeads, "A_log must have shape [48]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == kNumVHeads, "dt_bias must have shape [48]."); + TORCH_CHECK( + recurrent_state.dim() == 4 && + recurrent_state.size(1) == kNumVHeads && + recurrent_state.size(2) == kHeadDimQK && + recurrent_state.size(3) == kHeadDimV, + "recurrent_state must have shape [pool, 48, 128, 128]."); + TORCH_CHECK(pool_idx.dim() == 1 && pool_idx.size(0) == token_count, "pool_idx must have shape [N]."); + TORCH_CHECK( + out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimV}), + "out must have shape [N, 48, 128]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + q_rep.scalar_type(), + "launch_qwen35_scalar_kda_decode_kernel", + [&] { + kernel::launch_qwen35_scalar_kda_decode_kernel( + stream, + q_rep.data_ptr(), + k_rep.data_ptr(), + v.data_ptr(), + a_kernel.data_ptr(), + b_kernel.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + recurrent_state.data_ptr(), + pool_idx.data_ptr(), + out.data_ptr(), + static_cast(token_count)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp new file mode 100644 index 00000000..7628a3ce --- /dev/null +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -0,0 +1,247 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "qwen35_decode_common.cuh" +#include "qwen35_scalar_kda_mainloop.hpp" + +#include + +namespace cula::qwen35::decode::kernel { + +using namespace cute; + +template +struct Qwen35ScalarKdaDecodeKernel { + // Decode-first design: + // - 1 CTA owns 1 (token_idx, hv) + // - 1 warpgroup (128 threads) per CTA + // - recurrent state stays fp32 and is traversed as 16x16 tiles over the + // internal [V, K] view + // - the intended optimized path is fp32 FFMA on CUDA cores, not a forced + // Tensor Core lowering + static constexpr int kThreads = 128; + static constexpr int kWarpGroupThreads = 128; + static constexpr int kTileV = 16; + static constexpr int kTileK = 16; + static constexpr int kTilesPerV = kHeadDimV / kTileV; + static constexpr int kTilesPerK = kHeadDimQK / kTileK; + + static_assert(kNumQKHeads < kNumVHeads); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kThreads == kWarpGroupThreads); + static_assert(kHeadDimV % kTileV == 0); + static_assert(kHeadDimQK % kTileK == 0); + + struct SharedStorage { + // Shared staging plan for the fp32 decode path: + // - q/k/v are staged once per CTA + // - proj/out intermediates remain in fp32 + // - recurrent state itself remains in fp32 global storage + alignas(16) scalar_t q_smem[kHeadDimQK]; + alignas(16) scalar_t k_smem[kHeadDimQK]; + alignas(16) scalar_t v_smem[kHeadDimV]; + alignas(16) float proj_smem[kHeadDimV]; + alignas(16) float out_smem[kHeadDimV]; + }; + + static dim3 block_shape() { + return dim3(kThreads, 1, 1); + } + + static dim3 grid_shape(int token_count) { + // One block owns one (token_idx, hv) pair in the first implementation. + return dim3(static_cast(kNumVHeads), static_cast(token_count), 1); + } + + template + CUTE_DEVICE static void run_device( + const scalar_t* __restrict__ q_rep, + const scalar_t* __restrict__ k_rep, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a_kernel, + const scalar_t* __restrict__ b_kernel, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count, + SharedStorage& storage) { + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + if (token_idx >= token_count || hv >= kNumVHeads) { + return; + } + + // Internal tensor-view contract fixed for the first implementation pass: + // + // 1. q_rep / k_rep / v / out stay in their external contiguous layouts: + // - q_rep : [N, HV, K] with stride (HV*K, K, 1) + // - k_rep : [N, HV, K] with stride (HV*K, K, 1) + // - v : [N, HV, V] with stride (HV*V, V, 1) + // - out : [N, HV, V] with stride (HV*V, V, 1) + // + // 2. a_kernel / b_kernel are treated as: + // - [N, HV] with stride (HV, 1) + // + // 3. A_log / dt_bias are treated as: + // - [HV] with stride (1) + // + // 4. recurrent_state keeps the external physical storage contract: + // - [pool, HV, K, V] + // but the kernel's main computation will use an internal VK view: + // - [pool, HV, V, K] + // + // This lets the recurrent update consume one V-row of state against q/k + // more naturally in the first mainloop design, while preserving the + // existing external state ABI. + // + // The current block owns exactly one (token_idx, hv) pair. That means one + // warpgroup-sized CTA updates one 128x128 recurrent-state tile for one + // v-head. + // + // TODO(qwen35-scalar-kda-opt): + // - Likely next optimization path: keep one CTA per (token_idx, hv), but + // tile the 128x128 state more aggressively inside the block (for example + // along V tiles or KxV subtiles assigned per warp). + // - More complex alternative: split one (token_idx, hv) tile across + // multiple CTAs and coordinate updates. Not a first-pass target. + // - After the fp32 decode path is stable, evaluate warp specialization: + // dedicated producer/load warp(s) vs consumer/compute warp(s), instead + // of introducing that complexity before the math path itself is stable. + + auto q_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kNumVHeads * kHeadDimQK, kHeadDimQK, Int<1>{})); + auto v_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kNumVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kNumVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_kv = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + auto gQ = make_tensor(make_gmem_ptr(q_rep), q_layout); + auto gK = make_tensor(make_gmem_ptr(k_rep), q_layout); + auto gV = make_tensor(make_gmem_ptr(v), v_layout); + auto gO = make_tensor(make_gmem_ptr(out), v_layout); + auto gA = make_tensor(make_gmem_ptr(a_kernel), head_layout); + auto gB = make_tensor(make_gmem_ptr(b_kernel), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_kv = make_tensor(make_gmem_ptr(recurrent_state), state_layout_kv); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + (void)gH_kv; // Keep the physical KV view documented and available. + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, hv, _); + auto k_vec = gK(token_idx, hv, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto a_scalar = gA(token_idx, hv); + auto b_scalar = gB(token_idx, hv); + auto A_log_scalar = gAlog(hv); + auto dt_bias_scalar = gDt(hv); + auto state_vk = gH_vk(state_row, hv, _, _); + + Mainloop::run( + q_vec, + k_vec, + v_vec, + a_scalar, + b_scalar, + A_log_scalar, + dt_bias_scalar, + state_vk, + out_vec, + storage, + tid, + kThreads); + } +}; + +template > +__global__ void qwen35_scalar_kda_decode_kernel( + const scalar_t* __restrict__ q_rep, + const scalar_t* __restrict__ k_rep, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a_kernel, + const scalar_t* __restrict__ b_kernel, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; + Qwen35ScalarKdaDecodeKernel::template run_device( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count, + storage); +} + +template > +void launch_qwen35_scalar_kda_decode_kernel( + cudaStream_t stream, + const scalar_t* q_rep, + const scalar_t* k_rep, + const scalar_t* v, + const scalar_t* a_kernel, + const scalar_t* b_kernel, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_scalar_kda_decode_kernel<<>>( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + +} // namespace cula::qwen35::decode::kernel diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp new file mode 100644 index 00000000..454dd146 --- /dev/null +++ b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp @@ -0,0 +1,452 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "qwen35_decode_common.cuh" + +#include +#include +#include +#include + +namespace cula::qwen35::decode::kernel { + +using namespace cute; + +template +struct Qwen35ScalarKdaDecodeMainloop { + // Decode design decision: + // - recurrent_state remains fp32 both physically and mathematically + // - decode is treated as a register-level recurrent GEMV/rank-1-update + // problem, not as a Tensor Core GEMM problem + // - you can think of the target implementation style as a + // flash_linear_decode_kernel: pure CUDA Core math, warp-shuffle vector + // sharing, and fp32 state kept live as long as possible during one token + // update + // + // Reason: + // - state participates in a long recurrent chain; lowering master-state + // precision is risky and usually not worth it + // - decode operates on a single-token q/k vector, so the dominant kernels + // are GEMV-like: + // proj = state @ k + // out = state' @ q + // This is typically register / memory bound rather than Tensor-Core bound + // + // Practical consequence: + // - the first production-worthy decode path should be built around fp32 FFMA + // on CUDA cores + // - tile structure is still useful, but it should serve register ownership, + // reduction, and cache behavior instead of forcing a GMMA lowering + // + // The remaining work for this kernel is therefore: + // 1. tighten thread ownership of the fp32 state tile + // 2. optimize proj / update / out reductions + // 3. evaluate warp-specialized load/compute roles only after the fp32 path + // is stable and measured + static constexpr int kTileV = 16; + static constexpr int kTileK = 16; + static constexpr int kTilesPerV = kHeadDimV / kTileV; + static constexpr int kTilesPerK = kHeadDimQK / kTileK; + static constexpr int kWarpSize = 32; + static constexpr int kRowsPerTile = kTileV; + static constexpr int kWarpsPerCta = 4; + static constexpr int kRowsPerWarp = kWarpSize; + static constexpr int kRowsPerThread = 1; + + static_assert(kHeadDimV == 128); + static_assert(kHeadDimQK == 128); + static_assert(kWarpsPerCta * kRowsPerWarp == kHeadDimV); + + // First concrete decode threading plan: + // + // - 1 CTA = 1 (token, hv) + // - 128 threads = 4 warps + // - 1 thread owns exactly 1 V-row of the 128x128 recurrent state + // - Therefore one CTA covers all 128 V-rows exactly once + // + // For the owned row, the thread streams over K in 16-wide tiles: + // + // state_row[0:15] -> registers + // state_row[16:31] -> registers + // ... + // state_row[112:127]-> registers + // + // This means the first concrete fp32 path does NOT attempt to keep the + // whole 128-float row resident in registers at once. Instead it keeps the + // current K tile resident: + // + // - state_regs[16] : current fp32 state tile + // - k_regs[16] : current key tile + // - q_regs[16] : current query tile + // + // plus a handful of scalar accumulators: + // + // - proj_row + // - out_row + // - v_new_row + // - gate scalars + // + // This is a practical first step toward the user's desired "state stays in + // registers for the current token" behavior while keeping register pressure + // manageable. + // + // Reduction policy for this first concrete plan: + // + // - proj/out are row-local, so no warp reduction is required + // - each row is fully owned by one thread across all K tiles + // - warp shuffle is reserved for future vector-broadcast refinements if we + // decide to move q/k staging from shared memory into warp-register paths + struct ThreadRowPlan { + int warp_id; + int lane_id; + int v_row; + bool owns_row; + }; + + struct TileCoords { + int v_base; + int k_base; + }; + + CUTE_DEVICE static ThreadRowPlan make_thread_row_plan(int tid) { + const int warp_id = tid / kWarpSize; + const int lane_id = tid % kWarpSize; + const int v_row = warp_id * kRowsPerWarp + lane_id; + const bool owns_row = v_row < kHeadDimV; + return ThreadRowPlan{warp_id, lane_id, v_row, owns_row}; + } + + template + CUTE_DEVICE static void load_vec_tile_to_regs( + TensorVec const& vec, + TileCoords coords, + float (®s)[kTileK]) { +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + regs[kk] = static_cast(vec(coords.k_base + kk)); + } + } + + template + CUTE_DEVICE static void load_state_row_tile_to_regs( + TensorState const& state_vk, + int v_row, + TileCoords coords, + float (&state_regs)[kTileK]) { +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + state_regs[kk] = static_cast(state_vk(v_row, coords.k_base + kk)); + } + } + + template + CUTE_DEVICE static void store_state_row_tile_from_regs( + TensorState& state_vk, + int v_row, + TileCoords coords, + float const (&state_regs)[kTileK]) { +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + state_vk(v_row, coords.k_base + kk) = state_regs[kk]; + } + } + + struct RowTileProjPlan { + int v_base; + int k_base; + int warp_id; + int lane_id; + bool owns_row; + int row_in_tile; + int v_row; + }; + + CUTE_DEVICE static TileCoords make_tile_coords(int tile_v, int tile_k) { + return TileCoords{tile_v * kTileV, tile_k * kTileK}; + } + + CUTE_DEVICE static RowTileProjPlan make_row_tile_proj_plan( + TileCoords coords, + int warp_id, + int lane_id) { + const bool owns_row = lane_id < kTileV; + const int row_in_tile = lane_id; + const int v_row = coords.v_base + row_in_tile; + return RowTileProjPlan{ + coords.v_base, + coords.k_base, + warp_id, + lane_id, + owns_row, + row_in_tile, + v_row, + }; + } + + struct RowTileUpdatePlan { + TileCoords coords; + int warp_id; + int lane_id; + bool owns_row; + int row_in_tile; + int v_row; + }; + + CUTE_DEVICE static RowTileUpdatePlan make_row_tile_update_plan( + TileCoords coords, + int warp_id, + int lane_id) { + const bool owns_row = lane_id < kTileV; + const int row_in_tile = lane_id; + const int v_row = coords.v_base + row_in_tile; + return RowTileUpdatePlan{ + coords, + warp_id, + lane_id, + owns_row, + row_in_tile, + v_row, + }; + } + + template + CUTE_DEVICE static float accumulate_proj_row_tile( + TensorState const& state_vk, + TensorKTile const& k_smem, + int v_row, + TileCoords coords) { + float state_regs[kTileK]; + float k_regs[kTileK]; + load_state_row_tile_to_regs(state_vk, v_row, coords, state_regs); + load_vec_tile_to_regs(k_smem, coords, k_regs); + + float accum = 0.f; +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + accum += state_regs[kk] * k_regs[kk]; + } + return accum; + } + + template + CUTE_DEVICE static float update_state_row_tile_and_accumulate_out( + TensorState& state_vk, + TensorKTile const& k_smem, + TensorQTile const& q_smem, + int v_row, + TileCoords coords, + float decay, + float v_new) { + float state_regs[kTileK]; + float k_regs[kTileK]; + float q_regs[kTileK]; + load_state_row_tile_to_regs(state_vk, v_row, coords, state_regs); + load_vec_tile_to_regs(k_smem, coords, k_regs); + load_vec_tile_to_regs(q_smem, coords, q_regs); + + float out_acc = 0.f; +#pragma unroll + for (int kk = 0; kk < kTileK; ++kk) { + const float state_new = decay * state_regs[kk] + v_new * k_regs[kk]; + state_regs[kk] = state_new; + out_acc += state_new * q_regs[kk]; + } + store_state_row_tile_from_regs(state_vk, v_row, coords, state_regs); + return out_acc; + } + + template + CUTE_DEVICE static float project_row_tile( + TensorState const& state_vk, + TensorKTile const& k_smem, + int v_row, + TileCoords coords) { + // Current decode path: + // - one thread owns one full V-row + // - this helper computes the row-local proj contribution for one K tile + // - no cross-thread reduction is needed + return accumulate_proj_row_tile(state_vk, k_smem, v_row, coords); + } + + template + CUTE_DEVICE static float project_row_tile( + TensorState const& state_vk, + TensorKTile const& k_smem, + RowTileProjPlan const& plan) { + if (!plan.owns_row) { + return 0.f; + } + return project_row_tile( + state_vk, k_smem, plan.v_row, TileCoords{plan.v_base, plan.k_base}); + } + + template + CUTE_DEVICE static float update_and_output_row_tile( + TensorState& state_vk, + TensorKTile const& k_smem, + TensorQTile const& q_smem, + int v_row, + TileCoords coords, + float decay, + float v_new) { + // Current decode path: + // - read one 16-wide state tile for the owned row into registers + // - apply decay and rank-1 update in fp32 + // - accumulate the matching out contribution against q + // - write the updated state tile back + return update_state_row_tile_and_accumulate_out( + state_vk, k_smem, q_smem, v_row, coords, decay, v_new); + } + + template + CUTE_DEVICE static float update_and_output_row_tile( + TensorState& state_vk, + TensorKTile const& k_smem, + TensorQTile const& q_smem, + RowTileUpdatePlan const& plan, + float decay, + float v_new) { + if (!plan.owns_row) { + return 0.f; + } + return update_and_output_row_tile( + state_vk, + k_smem, + q_smem, + plan.v_row, + plan.coords, + decay, + v_new); + } + + CUTE_DEVICE static float softplusf_approx(float x) { + return x > 20.f ? x : log1pf(expf(x)); + } + + template < + typename TensorQ, + typename TensorK, + typename TensorV, + typename TensorA, + typename TensorB, + typename TensorAlog, + typename TensorDt, + typename TensorHvk, + typename TensorOut, + typename SharedStorage> + CUTE_DEVICE static void run( + TensorQ const& q_vec, + TensorK const& k_vec, + TensorV const& v_vec, + TensorA const& a_scalar, + TensorB const& b_scalar, + TensorAlog const& A_log_scalar, + TensorDt const& dt_bias_scalar, + TensorHvk& state_vk, + TensorOut& out_vec, + SharedStorage& storage, + int tid, + int num_threads) { + // Decode organization: + // - 1 warpgroup owns the full [128, 128] state tile for one (token, hv) + // - state is traversed as 16x16 tiles over the internal VK view + // - q/k/v are staged once into shared memory + // - proj/out are accumulated over K tiles + // - rank-1 update is applied tile-by-tile in the same traversal order + // + // This pass establishes the tile-first organization for the final fp32 + // decode kernel. The next implementation step should optimize the scalar + // inner loops with better register ownership / reductions rather than + // forcing Tensor Core math. + // + // TODO(qwen35-decode-fp32): + // - evaluate whether q/k should move from shared-memory staging to + // warp-shuffle broadcast + // - evaluate whether one thread should own more than one V-row + // - evaluate whether some parts of the state row can remain resident in + // registers across both proj and update/out passes with acceptable + // register pressure + + const float a_val = static_cast(a_scalar()); + const float b_val = static_cast(b_scalar()); + const float A_log_val = static_cast(A_log_scalar()); + const float dt_bias_val = static_cast(dt_bias_scalar()); + + const float g = -expf(A_log_val) * softplusf_approx(a_val + dt_bias_val); + const float decay = expf(g); + const float beta = 1.f / (1.f + expf(-b_val)); + + auto q_smem = make_tensor(make_smem_ptr(storage.q_smem), make_layout(make_shape(Int{}))); + auto k_smem = make_tensor(make_smem_ptr(storage.k_smem), make_layout(make_shape(Int{}))); + auto v_smem = make_tensor(make_smem_ptr(storage.v_smem), make_layout(make_shape(Int{}))); + auto proj_smem = make_tensor(make_smem_ptr(storage.proj_smem), make_layout(make_shape(Int{}))); + auto out_smem = make_tensor(make_smem_ptr(storage.out_smem), make_layout(make_shape(Int{}))); + + // Stage q/k/v once per CTA for the current decode token. + for (int idx = tid; idx < kHeadDimQK; idx += num_threads) { + q_smem(idx) = q_vec(idx); + k_smem(idx) = k_vec(idx); + } + for (int idx = tid; idx < kHeadDimV; idx += num_threads) { + v_smem(idx) = v_vec(idx); + proj_smem(idx) = 0.f; + out_smem(idx) = 0.f; + } + __syncthreads(); + + ThreadRowPlan row_plan = make_thread_row_plan(tid); + + // First concrete ownership model: + // - each thread owns one full state row across all 128 K columns + // - the row is streamed tile-by-tile through registers + // - no cross-thread reduction is needed for proj/out because the full row + // stays with one thread for the duration of the token update + if (row_plan.owns_row) { + float proj_row = 0.f; + for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { + TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; + RowTileProjPlan proj_plan = make_row_tile_proj_plan(coords, row_plan.warp_id, row_plan.lane_id); + proj_plan.v_row = row_plan.v_row; + proj_plan.owns_row = true; + proj_row += project_row_tile(state_vk, k_smem, proj_plan); + } + + proj_smem(row_plan.v_row) = proj_row; + + const float v_val = static_cast(v_smem(row_plan.v_row)); + const float v_new_row = beta * (v_val - proj_row); + + float out_row = 0.f; + for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { + TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; + RowTileUpdatePlan update_plan = make_row_tile_update_plan(coords, row_plan.warp_id, row_plan.lane_id); + update_plan.v_row = row_plan.v_row; + update_plan.owns_row = true; + out_row += update_and_output_row_tile( + state_vk, k_smem, q_smem, update_plan, decay, v_new_row); + } + + out_smem(row_plan.v_row) = out_row; + } + __syncthreads(); + + for (int idx = tid; idx < kHeadDimV; idx += num_threads) { + out_vec(idx) = static_cast(out_smem(idx)); + } + } +}; + +} // namespace cula::qwen35::decode::kernel diff --git a/cula/ops/qwen35_conv1d_decode.py b/cula/ops/qwen35_conv1d_decode.py new file mode 100644 index 00000000..2481bead --- /dev/null +++ b/cula/ops/qwen35_conv1d_decode.py @@ -0,0 +1,170 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL kernel for Qwen3.5 single-token conv-state update.""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +THREADS = 256 +KERNEL_SIZE = 4 + + +@cute.kernel +def _qwen35_conv1d_decode_kernel( + x_t: cute.Tensor, + conv_state: cute.Tensor, + weight: cute.Tensor, + y: cute.Tensor, + B: cutlass.Constexpr[int], + C: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + linear_idx = bidx * THREADS + tidx + if linear_idx < B * C: + b = linear_idx // C + c = linear_idx % C + + s0 = cutlass.Float32(conv_state[(b, c, 1)]) + s1 = cutlass.Float32(conv_state[(b, c, 2)]) + s2 = cutlass.Float32(conv_state[(b, c, 3)]) + s3 = cutlass.Float32(x_t[(b, c)]) + + w0 = cutlass.Float32(weight[(c, 0)]) + w1 = cutlass.Float32(weight[(c, 1)]) + w2 = cutlass.Float32(weight[(c, 2)]) + w3 = cutlass.Float32(weight[(c, 3)]) + + out = s0 * w0 + s1 * w1 + s2 * w2 + s3 * w3 + sig = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-out)) + out = out * sig + + conv_state[(b, c, 0)] = cutlass.BFloat16(s0) + conv_state[(b, c, 1)] = cutlass.BFloat16(s1) + conv_state[(b, c, 2)] = cutlass.BFloat16(s2) + conv_state[(b, c, 3)] = cutlass.BFloat16(s3) + y[(b, c)] = cutlass.BFloat16(out) + + +@cute.jit +def _run_qwen35_conv1d_decode( + x_t: cute.Tensor, + conv_state: cute.Tensor, + weight: cute.Tensor, + y: cute.Tensor, + B: cutlass.Constexpr[int], + C: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + stream: cuda.CUstream, +): + _qwen35_conv1d_decode_kernel( + x_t, + conv_state, + weight, + y, + B, + C, + K, + ).launch( + grid=(cute.ceil_div(B * C, THREADS), 1, 1), + block=(THREADS, 1, 1), + stream=stream, + ) + + +@functools.cache +def _get_compiled_kernel( + B: int, + C: int, + K: int, +): + return {} + + +def qwen35_conv1d_decode_update( + x_t: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + *, + activation: str = "silu", +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-token depthwise causal conv1d update. + + Expected shapes: + - x_t: [B, C] + - conv_state: [B, C, 4] + - weight: [C, 1, 4] or [C, 4] + """ + + if activation != "silu": + raise ValueError(f"Only silu activation is currently supported, got {activation}") + if x_t.ndim != 2: + raise ValueError(f"x_t must be 2D [batch, channels], got {tuple(x_t.shape)}") + if conv_state.ndim != 3: + raise ValueError(f"conv_state must be 3D [batch, channels, kernel], got {tuple(conv_state.shape)}") + if conv_state.shape[:2] != x_t.shape: + raise ValueError(f"conv_state batch/channel dims must match x_t, got x_t={tuple(x_t.shape)} conv_state={tuple(conv_state.shape)}") + kernel_size = conv_state.shape[-1] + if kernel_size != 4: + raise ValueError(f"Expected kernel_size=4 for Qwen3.5, got {kernel_size}") + + if weight.ndim == 3: + if weight.shape[1] != 1 or weight.shape[2] != kernel_size: + raise ValueError(f"weight must be [channels,1,{kernel_size}], got {tuple(weight.shape)}") + weight_2d = weight.squeeze(1) + elif weight.ndim == 2: + if weight.shape[1] != kernel_size: + raise ValueError(f"weight must be [channels,{kernel_size}], got {tuple(weight.shape)}") + weight_2d = weight + else: + raise ValueError(f"weight must be 2D or 3D, got {tuple(weight.shape)}") + + if weight_2d.shape[0] != x_t.shape[1]: + raise ValueError(f"weight channels must match x_t channels, got weight={tuple(weight_2d.shape)} x_t={tuple(x_t.shape)}") + + x_t = x_t.contiguous() + conv_state = conv_state.contiguous() + weight_2d = weight_2d.contiguous() + y = torch.empty_like(x_t) + + B, C = x_t.shape + cache = _get_compiled_kernel(B, C, kernel_size) + if "compiled" not in cache: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + compiled = cute.compile( + _run_qwen35_conv1d_decode, + from_dlpack(x_t, assumed_align=16), + from_dlpack(conv_state, assumed_align=16), + from_dlpack(weight_2d, assumed_align=16), + from_dlpack(y, assumed_align=16), + B=B, + C=C, + K=kernel_size, + stream=stream, + options="--enable-tvm-ffi", + ) + cache["compiled"] = compiled + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + cache["compiled"](x_t, conv_state, weight_2d, y, stream) + return y, conv_state diff --git a/cula/ops/qwen35_conv1d_prefill.py b/cula/ops/qwen35_conv1d_prefill.py new file mode 100644 index 00000000..e153460b --- /dev/null +++ b/cula/ops/qwen35_conv1d_prefill.py @@ -0,0 +1,36 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL placeholder for Qwen3.5 depthwise causal conv1d prefill.""" + +from __future__ import annotations + +import torch + + +def qwen35_conv1d_prefill( + x: torch.Tensor, + weight: torch.Tensor, + *, + activation: str = "silu", +) -> torch.Tensor: + """Depthwise causal conv1d over a full sequence. + + Expected shapes: + - x: [B, C, S] + - weight: [C, 1, 4] or [C, 4] + """ + + del x, weight, activation + raise NotImplementedError("Qwen3.5 conv1d prefill kernel is not implemented yet.") diff --git a/cula/ops/qwen35_scalar_kda_decode.py b/cula/ops/qwen35_scalar_kda_decode.py new file mode 100644 index 00000000..db6ece66 --- /dev/null +++ b/cula/ops/qwen35_scalar_kda_decode.py @@ -0,0 +1,75 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL placeholder for Qwen3.5 scalar-gated KDA decode.""" + +from __future__ import annotations + +import torch + +from cula.ops.kda_decode import kda_decode + + +def qwen35_scalar_kda_decode( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + *, + state_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-token scalar-gated delta-rule decode for Qwen3.5.""" + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError(f"q/k/v must be 4D, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + if q.shape != k.shape: + raise ValueError(f"q and k must have the same shape, got q={tuple(q.shape)} vs k={tuple(k.shape)}") + if q.shape[1] != 1 or v.shape[1] != 1: + raise ValueError(f"Decode expects single-token sequence dim, got q={tuple(q.shape)} v={tuple(v.shape)}") + + N, _, HV, K = q.shape + if a.ndim == 2: + a = a.unsqueeze(1) + if b.ndim == 2: + b = b.unsqueeze(1) + if a.shape != (N, 1, HV) or b.shape != (N, 1, HV): + raise ValueError(f"a/b must be [N,1,HV], got a={tuple(a.shape)} b={tuple(b.shape)}") + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + + a_expanded = a.unsqueeze(-1).expand(N, 1, HV, K) + dt_bias_expanded = dt_bias[:, None].expand(HV, K).contiguous() + state_indices = ( + torch.arange(N, device=q.device, dtype=torch.int32) + if state_indices is None + else state_indices.to(device=q.device, dtype=torch.int32) + ) + o = kda_decode( + A_log=A_log.contiguous(), + dt_bias=dt_bias_expanded, + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + a=a_expanded.contiguous(), + b=b.contiguous(), + initial_state_source=recurrent_state, + initial_state_indices=state_indices, + scale=K**-0.5, + use_qk_l2norm_in_kernel=True, + state_layout="kv", + ) + return o, recurrent_state diff --git a/cula/ops/qwen35_scalar_kda_prefill.py b/cula/ops/qwen35_scalar_kda_prefill.py new file mode 100644 index 00000000..2fda06cb --- /dev/null +++ b/cula/ops/qwen35_scalar_kda_prefill.py @@ -0,0 +1,37 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CuTe DSL placeholder for Qwen3.5 scalar-gated KDA prefill.""" + +from __future__ import annotations + +import torch + + +def qwen35_scalar_kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Chunked scalar-gated delta-rule prefill for Qwen3.5.""" + + del q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens + raise NotImplementedError("Qwen3.5 scalar-gated KDA prefill kernel is not implemented yet.") diff --git a/cula/qwen35/__init__.py b/cula/qwen35/__init__.py new file mode 100644 index 00000000..a8fd8351 --- /dev/null +++ b/cula/qwen35/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3.5-specific linear attention support built on top of cuLA primitives.""" + +from cula.qwen35.common import Qwen35LinearAttentionConfig +from cula.qwen35.runtime import ( + qwen35_linear_attention_decode, + qwen35_linear_attention_prefill, +) + +__all__ = [ + "Qwen35LinearAttentionConfig", + "qwen35_linear_attention_prefill", + "qwen35_linear_attention_decode", +] diff --git a/cula/qwen35/common.py b/cula/qwen35/common.py new file mode 100644 index 00000000..63299c06 --- /dev/null +++ b/cula/qwen35/common.py @@ -0,0 +1,135 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared constants and validation helpers for Qwen3.5 linear attention.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class Qwen35LinearAttentionConfig: + """Minimal runtime config for Qwen3.5 linear-attention kernels.""" + + hidden_size: int = 5120 + conv_kernel_size: int = 4 + num_k_heads: int = 16 + num_v_heads: int = 48 + head_k_dim: int = 128 + head_v_dim: int = 128 + qkv_dtype: torch.dtype = torch.bfloat16 + state_dtype: torch.dtype = torch.float32 + + @property + def key_dim(self) -> int: + return self.num_k_heads * self.head_k_dim + + @property + def value_dim(self) -> int: + return self.num_v_heads * self.head_v_dim + + @property + def conv_dim(self) -> int: + return self.key_dim * 2 + self.value_dim + + @property + def qk_repeat_factor(self) -> int: + assert self.num_v_heads % self.num_k_heads == 0 + return self.num_v_heads // self.num_k_heads + + +DEFAULT_QWEN35_LINEAR_ATTN_CONFIG = Qwen35LinearAttentionConfig() + + +def validate_mixed_qkv( + mixed_qkv: torch.Tensor, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> None: + if mixed_qkv.dtype != config.qkv_dtype: + raise TypeError(f"mixed_qkv must be {config.qkv_dtype}, got {mixed_qkv.dtype}") + if mixed_qkv.ndim != 2: + raise ValueError(f"mixed_qkv must be 2D [tokens, conv_dim_local], got {tuple(mixed_qkv.shape)}") + if mixed_qkv.shape[-1] <= 0: + raise ValueError("mixed_qkv must have a non-zero channel dimension") + if mixed_qkv.shape[-1] % config.conv_dim != 0 and mixed_qkv.shape[-1] != config.conv_dim: + # In TP mode this is expected to be a local shard, so only require alignment + # with the Qwen3.5 packed layout ratio. + local_dim = mixed_qkv.shape[-1] + expected_splits = (config.key_dim, config.key_dim, config.value_dim) + if local_dim % sum(expected_splits) != 0: + raise ValueError(f"mixed_qkv last dim must match packed local conv dim, got {local_dim}") + + +def validate_scalar_gate_inputs( + a: torch.Tensor, + b: torch.Tensor, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> None: + if a.shape != b.shape: + raise ValueError(f"a and b must have the same shape, got a={tuple(a.shape)} vs b={tuple(b.shape)}") + if a.ndim != 2: + raise ValueError(f"a and b must be 2D [tokens, num_v_heads_local], got {tuple(a.shape)}") + if a.dtype != config.qkv_dtype or b.dtype != config.qkv_dtype: + raise TypeError(f"a and b must be {config.qkv_dtype}, got a={a.dtype}, b={b.dtype}") + + +def validate_state_tensors( + conv_state: torch.Tensor | None, + recurrent_state: torch.Tensor | None, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> None: + if conv_state is not None: + if conv_state.ndim != 3: + raise ValueError(f"conv_state must be 3D [batch, channels, {config.conv_kernel_size}], got {tuple(conv_state.shape)}") + if conv_state.dtype != config.qkv_dtype: + raise TypeError(f"conv_state must be {config.qkv_dtype}, got {conv_state.dtype}") + if recurrent_state is not None: + if recurrent_state.ndim != 4: + raise ValueError(f"recurrent_state must be 4D [batch, hv, k, v], got {tuple(recurrent_state.shape)}") + if recurrent_state.dtype != config.state_dtype: + raise TypeError(f"recurrent_state must be {config.state_dtype}, got {recurrent_state.dtype}") + + +def infer_local_config( + mixed_qkv_dim: int, + local_num_v_heads: int, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> tuple[int, int, int]: + """Infer local packed dims from runtime shard sizes. + + Returns: + - local_key_dim + - local_value_dim + - local_num_k_heads + """ + + local_value_dim = local_num_v_heads * config.head_v_dim + remaining = mixed_qkv_dim - local_value_dim + if remaining <= 0 or remaining % 2 != 0: + raise ValueError( + f"Cannot infer local q/k dims from mixed_qkv_dim={mixed_qkv_dim}, local_num_v_heads={local_num_v_heads}" + ) + local_key_dim = remaining // 2 + if local_key_dim % config.head_k_dim != 0: + raise ValueError(f"Local key dim must be divisible by head_k_dim={config.head_k_dim}, got {local_key_dim}") + local_num_k_heads = local_key_dim // config.head_k_dim + if local_num_v_heads % local_num_k_heads != 0: + raise ValueError( + f"Local num_v_heads={local_num_v_heads} must be divisible by local num_k_heads={local_num_k_heads}" + ) + return local_key_dim, local_value_dim, local_num_k_heads diff --git a/cula/qwen35/runtime.py b/cula/qwen35/runtime.py new file mode 100644 index 00000000..47f9548a --- /dev/null +++ b/cula/qwen35/runtime.py @@ -0,0 +1,165 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime dispatch for Qwen3.5 linear-attention kernels.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import torch + +from cula.qwen35.common import ( + DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + Qwen35LinearAttentionConfig, + infer_local_config, + validate_mixed_qkv, + validate_scalar_gate_inputs, + validate_state_tensors, +) +from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_update +from cula.ops.qwen35_scalar_kda_decode import qwen35_scalar_kda_decode + +_stream_cache: dict[tuple[str, int], cuda.CUstream] = {} + + +def _get_cached_stream(device: torch.device) -> cuda.CUstream: + stream_id = int(torch.cuda.current_stream(device=device).cuda_stream) + cache_key = (str(device), stream_id) + if cache_key not in _stream_cache: + _stream_cache[cache_key] = cuda.CUstream(stream_id) + return _stream_cache[cache_key] + + +def qwen35_linear_attention_prefill( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + cu_seqlens: torch.Tensor | None = None, + recurrent_state: torch.Tensor | None = None, + conv_state: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Qwen3.5 prefill wrapper. + + This is a thin runtime boundary. The underlying CuTe kernels are added in + dedicated `cula.ops.qwen35_*` modules. + """ + + del conv_weight, A_log, dt_bias, cu_seqlens + validate_mixed_qkv(mixed_qkv, config) + validate_scalar_gate_inputs(a, b, config) + validate_state_tensors(conv_state, recurrent_state, config) + _get_cached_stream(mixed_qkv.device) + raise NotImplementedError("Qwen3.5 prefill kernel path is not implemented yet.") + + +def qwen35_linear_attention_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + conv_state: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Qwen3.5 decode wrapper. + + Args: + mixed_qkv: [tokens, local_conv_dim] + a, b: [tokens, local_num_v_heads] + conv_weight: [local_conv_dim, 1, 4] or [local_conv_dim, 4] + A_log, dt_bias: [local_num_v_heads] + conv_state: [tokens, local_conv_dim, 4] + recurrent_state: [pool, local_num_v_heads, 128, 128] + + Returns: + - core_attn_out_flat: [tokens, local_value_dim] + - updated_conv_state + - updated_recurrent_state + """ + + validate_mixed_qkv(mixed_qkv, config) + validate_scalar_gate_inputs(a, b, config) + validate_state_tensors(conv_state, recurrent_state, config) + _get_cached_stream(mixed_qkv.device) + + if mixed_qkv.shape[0] != a.shape[0]: + raise ValueError(f"Token dimension mismatch, got mixed_qkv={tuple(mixed_qkv.shape)} a={tuple(a.shape)}") + if A_log.ndim != 1 or dt_bias.ndim != 1: + raise ValueError(f"A_log and dt_bias must be 1D, got {tuple(A_log.shape)} and {tuple(dt_bias.shape)}") + if A_log.shape != dt_bias.shape: + raise ValueError(f"A_log and dt_bias must have the same shape, got {tuple(A_log.shape)} vs {tuple(dt_bias.shape)}") + + tokens = mixed_qkv.shape[0] + local_num_v_heads = a.shape[1] + local_key_dim, local_value_dim, local_num_k_heads = infer_local_config( + mixed_qkv.shape[1], + local_num_v_heads, + config=config, + ) + if conv_state.shape != (tokens, mixed_qkv.shape[1], config.conv_kernel_size): + raise ValueError( + f"conv_state must be [tokens, local_conv_dim, {config.conv_kernel_size}], got {tuple(conv_state.shape)}" + ) + if recurrent_state.shape[1:] != (local_num_v_heads, config.head_k_dim, config.head_v_dim): + raise ValueError( + "recurrent_state must be [pool, local_num_v_heads, head_k_dim, head_v_dim], " + f"got {tuple(recurrent_state.shape)}" + ) + if A_log.numel() != local_num_v_heads: + raise ValueError(f"A_log must match local_num_v_heads={local_num_v_heads}, got {A_log.numel()}") + + conv_out, conv_state_out = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + ) + + q_end = local_key_dim + k_end = q_end + local_key_dim + q_flat = conv_out[:, :q_end] + k_flat = conv_out[:, q_end:k_end] + v_flat = conv_out[:, k_end:] + + q = q_flat.view(tokens, local_num_k_heads, config.head_k_dim) + k = k_flat.view(tokens, local_num_k_heads, config.head_k_dim) + v = v_flat.view(tokens, local_num_v_heads, config.head_v_dim) + + repeat_factor = local_num_v_heads // local_num_k_heads + q = q.repeat_interleave(repeat_factor, dim=1).unsqueeze(1).contiguous() + k = k.repeat_interleave(repeat_factor, dim=1).unsqueeze(1).contiguous() + v = v.unsqueeze(1).contiguous() + + core_attn_out, recurrent_state_out = qwen35_scalar_kda_decode( + q=q, + k=k, + v=v, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + ) + core_attn_out = core_attn_out.reshape(tokens, local_value_dim) + return core_attn_out, conv_state_out, recurrent_state_out diff --git a/docs/qwen35_kernel_plan.md b/docs/qwen35_kernel_plan.md new file mode 100644 index 00000000..5aeaba4d --- /dev/null +++ b/docs/qwen35_kernel_plan.md @@ -0,0 +1,40 @@ +# Qwen3.5 Kernel Landing Plan Inside cuLA + +This note records the internal landing structure for Qwen3.5 linear-attention +support added directly inside `cuLA`. + +## New Python package surface + +- `cula/qwen35/__init__.py` +- `cula/qwen35/common.py` +- `cula/qwen35/runtime.py` + +## New CuTe op entry files + +- `cula/ops/qwen35_conv1d_prefill.py` +- `cula/ops/qwen35_conv1d_decode.py` +- `cula/ops/qwen35_scalar_kda_prefill.py` +- `cula/ops/qwen35_scalar_kda_decode.py` + +## Intended ownership + +- `common.py` + shared constants, local-head config, shape validation +- `runtime.py` + compile-cache, stream-cache, prefill/decode dispatch boundaries +- `qwen35_conv1d_*` + depthwise causal conv1d + silu +- `qwen35_scalar_kda_*` + scalar-gated delta-rule prefill/decode kernels + +## What should be reused from existing cuLA code + +- runtime compile-cache patterns from `cula/ops/kda_decode.py` +- device helpers from `cula/utils.py` +- operator boundary style from `cula/kda/chunk.py` + +## What should stay isolated at first + +- no direct mutation of the generic `chunk_kda` public entry +- no pybind work until Python/CuTe path is numerically correct +- no conv + kda fusion until standalone kernels are validated From 7830fad02bb020ef479ce826c41b1b1ca6e64af9 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Fri, 5 Jun 2026 19:16:51 +0800 Subject: [PATCH 02/35] Wire Qwen3.5 decode runtime and tests --- csrc/api/pybind.cu | 70 +++++++++ csrc/qwen35/decode/qwen35_conv1d_decode.cu | 167 ++++++++++++++++++++- cula/__init__.py | 9 +- cula/ops/__init__.py | 18 ++- cula/ops/qwen35_conv1d_decode.py | 153 +++++++------------ cula/ops/qwen35_layout_decode.py | 111 ++++++++++++++ cula/ops/qwen35_scalar_kda_decode.py | 47 +++++- cula/qwen35/__init__.py | 19 +-- cula/qwen35/runtime.py | 156 ++++++++++++++++--- tests/test_qwen35_decode.py | 167 +++++++++++++++++++++ 10 files changed, 768 insertions(+), 149 deletions(-) create mode 100644 cula/ops/qwen35_layout_decode.py create mode 100644 tests/test_qwen35_decode.py diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index d14a41c5..cd05e5c0 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -17,6 +17,8 @@ #include #include +#include "qwen35/decode/qwen35_decode_common.cuh" + #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) void ChunkKDAFwdIntra( @@ -68,6 +70,71 @@ kda_fwd_prefill( bool safe_gate); #endif +void +qwen35_conv1d_decode( + at::Tensor mixed_qkv, + at::Tensor conv_state, + at::Tensor conv_weight, + at::Tensor out) { + cula::qwen35::decode::ConvDecodeParams params{ + mixed_qkv, + conv_state, + conv_weight, + out, + }; + cula::qwen35::decode::run_qwen35_conv1d_decode(params); +} + +void +qwen35_layout_decode( + at::Tensor mixed_qkv_conv, + at::Tensor a, + at::Tensor b, + at::Tensor q_rep, + at::Tensor k_rep, + at::Tensor v, + at::Tensor a_kernel, + at::Tensor b_kernel) { + cula::qwen35::decode::LayoutDecodeParams params{ + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + }; + cula::qwen35::decode::run_qwen35_layout_decode(params); +} + +void +qwen35_scalar_kda_decode( + at::Tensor q_rep, + at::Tensor k_rep, + at::Tensor v, + at::Tensor a_kernel, + at::Tensor b_kernel, + at::Tensor A_log, + at::Tensor dt_bias, + at::Tensor recurrent_state, + at::Tensor pool_idx, + at::Tensor out) { + cula::qwen35::decode::ScalarKdaDecodeParams params{ + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + }; + cula::qwen35::decode::run_qwen35_scalar_kda_decode(params); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "cuLA"; #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) @@ -77,4 +144,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #if defined(CULA_SM90A_ENABLED) m.def("kda_fwd_prefill", &kda_fwd_prefill); #endif + m.def("qwen35_conv1d_decode", &qwen35_conv1d_decode); + m.def("qwen35_layout_decode", &qwen35_layout_decode); + m.def("qwen35_scalar_kda_decode", &qwen35_scalar_kda_decode); } diff --git a/csrc/qwen35/decode/qwen35_conv1d_decode.cu b/csrc/qwen35/decode/qwen35_conv1d_decode.cu index cbce4896..34020375 100644 --- a/csrc/qwen35/decode/qwen35_conv1d_decode.cu +++ b/csrc/qwen35/decode/qwen35_conv1d_decode.cu @@ -14,16 +14,175 @@ #include "qwen35_decode_common.cuh" +#include +#include #include +#include +#include +#include + +namespace { + +template +__device__ inline float to_float(T x) { + return static_cast(x); +} + +template <> +__device__ inline float to_float(c10::Half x) { + return __half2float(static_cast<__half>(x)); +} + +template <> +__device__ inline float to_float(c10::BFloat16 x) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return __bfloat162float(static_cast<__nv_bfloat16>(x)); +#else + return static_cast(x); +#endif +} + +template +__device__ inline T from_float(float x) { + return static_cast(x); +} + +template <> +__device__ inline c10::Half from_float(float x) { + return c10::Half(__float2half_rn(x)); +} + +template <> +__device__ inline c10::BFloat16 from_float(float x) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return c10::BFloat16(__float2bfloat16(x)); +#else + return c10::BFloat16(x); +#endif +} + +template +__global__ void qwen35_conv1d_decode_kernel( + const scalar_t* __restrict__ mixed_qkv, + scalar_t* __restrict__ conv_state, + const scalar_t* __restrict__ conv_weight, + scalar_t* __restrict__ out, + int batch_size) { + constexpr int kThreads = 256; + const int64_t linear_idx = static_cast(blockIdx.x) * kThreads + threadIdx.x; + const int64_t total = static_cast(batch_size) * cula::qwen35::decode::kMixedQKVDim; + if (linear_idx >= total) { + return; + } + + const int64_t b = linear_idx / cula::qwen35::decode::kMixedQKVDim; + const int64_t c = linear_idx % cula::qwen35::decode::kMixedQKVDim; + + const int64_t x_idx = b * cula::qwen35::decode::kMixedQKVDim + c; + const int64_t state_base = + (b * cula::qwen35::decode::kMixedQKVDim + c) * cula::qwen35::decode::kConvKernelSize; + const int64_t weight_base = c * cula::qwen35::decode::kConvKernelSize; + + const float s0 = to_float(conv_state[state_base + 1]); + const float s1 = to_float(conv_state[state_base + 2]); + const float s2 = to_float(conv_state[state_base + 3]); + const float s3 = to_float(mixed_qkv[x_idx]); + + const float w0 = to_float(conv_weight[weight_base + 0]); + const float w1 = to_float(conv_weight[weight_base + 1]); + const float w2 = to_float(conv_weight[weight_base + 2]); + const float w3 = to_float(conv_weight[weight_base + 3]); + + const float conv = s0 * w0 + s1 * w1 + s2 * w2 + s3 * w3; + const float silu = conv / (1.f + expf(-conv)); + + conv_state[state_base + 0] = from_float(s0); + conv_state[state_base + 1] = from_float(s1); + conv_state[state_base + 2] = from_float(s2); + conv_state[state_base + 3] = from_float(s3); + out[x_idx] = from_float(silu); +} + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +} // namespace namespace cula::qwen35::decode { void run_qwen35_conv1d_decode(ConvDecodeParams& params) { - (void)params; + const at::Tensor& mixed_qkv = params.mixed_qkv; + const at::Tensor& conv_state = params.conv_state; + const at::Tensor& conv_weight = params.conv_weight; + const at::Tensor& out = params.out; + + TORCH_CHECK(mixed_qkv.is_cuda(), "mixed_qkv must be a CUDA tensor."); + const at::Device device = mixed_qkv.device(); + + check_tensor_device(conv_state, "conv_state", device); + check_tensor_device(conv_weight, "conv_weight", device); + check_tensor_device(out, "out", device); + + TORCH_CHECK(mixed_qkv.is_contiguous(), "mixed_qkv must be contiguous."); + TORCH_CHECK(conv_state.is_contiguous(), "conv_state must be contiguous."); + TORCH_CHECK(conv_weight.is_contiguous(), "conv_weight must be contiguous."); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous."); + + TORCH_CHECK( + mixed_qkv.scalar_type() == conv_state.scalar_type() && + mixed_qkv.scalar_type() == conv_weight.scalar_type() && + mixed_qkv.scalar_type() == out.scalar_type(), + "mixed_qkv/conv_state/conv_weight/out must share the same dtype."); + TORCH_CHECK( - false, - "run_qwen35_conv1d_decode is not implemented yet. " - "Planned kernel: single-token depthwise causal conv1d + silu for 10240 channels."); + mixed_qkv.scalar_type() == at::kHalf || mixed_qkv.scalar_type() == at::kBFloat16, + "conv decode only supports half/bfloat16."); + + const int64_t batch_size = mixed_qkv.size(0); + TORCH_CHECK( + mixed_qkv.dim() == 3 && mixed_qkv.sizes() == at::IntArrayRef({batch_size, 1, kMixedQKVDim}), + "mixed_qkv must have shape [B, 1, 10240]."); + TORCH_CHECK( + conv_state.dim() == 3 && + conv_state.sizes() == at::IntArrayRef({batch_size, kMixedQKVDim, kConvKernelSize}), + "conv_state must have shape [B, 10240, 4]."); + TORCH_CHECK( + (conv_weight.dim() == 2 && conv_weight.sizes() == at::IntArrayRef({kMixedQKVDim, kConvKernelSize})) || + (conv_weight.dim() == 3 && + conv_weight.sizes() == at::IntArrayRef({kMixedQKVDim, 1, kConvKernelSize})), + "conv_weight must have shape [10240, 4] or [10240, 1, 4]."); + TORCH_CHECK( + out.dim() == 3 && out.sizes() == at::IntArrayRef({batch_size, 1, kMixedQKVDim}), + "out must have shape [B, 1, 10240]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + const at::Tensor mixed_qkv_2d = mixed_qkv.view({batch_size, kMixedQKVDim}); + const at::Tensor out_2d = out.view({batch_size, kMixedQKVDim}); + const at::Tensor weight_2d = + conv_weight.dim() == 3 ? conv_weight.view({kMixedQKVDim, kConvKernelSize}) : conv_weight; + + constexpr int kThreads = 256; + const int64_t total = batch_size * static_cast(kMixedQKVDim); + const dim3 block(kThreads, 1, 1); + const dim3 grid(static_cast((total + kThreads - 1) / kThreads), 1, 1); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + mixed_qkv.scalar_type(), + "qwen35_conv1d_decode_kernel", + [&] { + qwen35_conv1d_decode_kernel<<>>( + mixed_qkv_2d.data_ptr(), + conv_state.data_ptr(), + weight_2d.data_ptr(), + out_2d.data_ptr(), + static_cast(batch_size)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } } // namespace cula::qwen35::decode diff --git a/cula/__init__.py b/cula/__init__.py index 7272e289..2688dab2 100644 --- a/cula/__init__.py +++ b/cula/__init__.py @@ -14,8 +14,9 @@ __version__ = "0.1.0" -from cula.ops.lightning_attn_sm100 import LinearAttentionChunkwiseDecay +try: + from cula.ops.lightning_attn_sm100 import LinearAttentionChunkwiseDecay +except Exception: # pragma: no cover - optional runtime dependency + LinearAttentionChunkwiseDecay = None -__all__ = [ - "LinearAttentionChunkwiseDecay", -] +__all__ = ["LinearAttentionChunkwiseDecay"] diff --git a/cula/ops/__init__.py b/cula/ops/__init__.py index 6450488b..99d34e0d 100644 --- a/cula/ops/__init__.py +++ b/cula/ops/__init__.py @@ -12,11 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode -from cula.ops.la_decode import linear_attention_decode +try: + from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode +except Exception: # pragma: no cover - optional runtime dependency + fused_sigmoid_gating_delta_rule_update = None + kda_decode = None -__all__ = [ - "kda_decode", - "fused_sigmoid_gating_delta_rule_update", - "linear_attention_decode", -] +try: + from cula.ops.la_decode import linear_attention_decode +except Exception: # pragma: no cover - optional runtime dependency + linear_attention_decode = None + +__all__ = ["kda_decode", "fused_sigmoid_gating_delta_rule_update", "linear_attention_decode"] diff --git a/cula/ops/qwen35_conv1d_decode.py b/cula/ops/qwen35_conv1d_decode.py index 2481bead..bd0af8f2 100644 --- a/cula/ops/qwen35_conv1d_decode.py +++ b/cula/ops/qwen35_conv1d_decode.py @@ -12,93 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CuTe DSL kernel for Qwen3.5 single-token conv-state update.""" +"""Qwen3.5 single-token conv-state update wrapper.""" from __future__ import annotations -import functools - -import cuda.bindings.driver as cuda -import cutlass -import cutlass.cute as cute import torch -from cutlass.cute.runtime import from_dlpack - -THREADS = 256 -KERNEL_SIZE = 4 - - -@cute.kernel -def _qwen35_conv1d_decode_kernel( - x_t: cute.Tensor, - conv_state: cute.Tensor, - weight: cute.Tensor, - y: cute.Tensor, - B: cutlass.Constexpr[int], - C: cutlass.Constexpr[int], - K: cutlass.Constexpr[int], -): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - linear_idx = bidx * THREADS + tidx - if linear_idx < B * C: - b = linear_idx // C - c = linear_idx % C - - s0 = cutlass.Float32(conv_state[(b, c, 1)]) - s1 = cutlass.Float32(conv_state[(b, c, 2)]) - s2 = cutlass.Float32(conv_state[(b, c, 3)]) - s3 = cutlass.Float32(x_t[(b, c)]) - - w0 = cutlass.Float32(weight[(c, 0)]) - w1 = cutlass.Float32(weight[(c, 1)]) - w2 = cutlass.Float32(weight[(c, 2)]) - w3 = cutlass.Float32(weight[(c, 3)]) - - out = s0 * w0 + s1 * w1 + s2 * w2 + s3 * w3 - sig = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-out)) - out = out * sig - - conv_state[(b, c, 0)] = cutlass.BFloat16(s0) - conv_state[(b, c, 1)] = cutlass.BFloat16(s1) - conv_state[(b, c, 2)] = cutlass.BFloat16(s2) - conv_state[(b, c, 3)] = cutlass.BFloat16(s3) - y[(b, c)] = cutlass.BFloat16(out) - - -@cute.jit -def _run_qwen35_conv1d_decode( - x_t: cute.Tensor, - conv_state: cute.Tensor, - weight: cute.Tensor, - y: cute.Tensor, - B: cutlass.Constexpr[int], - C: cutlass.Constexpr[int], - K: cutlass.Constexpr[int], - stream: cuda.CUstream, -): - _qwen35_conv1d_decode_kernel( - x_t, - conv_state, - weight, - y, - B, - C, - K, - ).launch( - grid=(cute.ceil_div(B * C, THREADS), 1, 1), - block=(THREADS, 1, 1), - stream=stream, - ) +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def qwen35_conv1d_decode_reference( + x_t: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pure torch reference for Qwen3.5 single-token depthwise conv decode.""" + if weight.ndim == 3: + weight = weight.squeeze(1) + + state_tail = conv_state[..., 1:].to(torch.float32) + x_last = x_t.unsqueeze(-1).to(torch.float32) + window = torch.cat([state_tail, x_last], dim=-1) + conv = (window * weight.to(torch.float32).unsqueeze(0)).sum(dim=-1) + y = torch.nn.functional.silu(conv).to(dtype=x_t.dtype) -@functools.cache -def _get_compiled_kernel( - B: int, - C: int, - K: int, -): - return {} + conv_state_out = conv_state.clone() + conv_state_out[..., 0] = conv_state[..., 1] + conv_state_out[..., 1] = conv_state[..., 2] + conv_state_out[..., 2] = conv_state[..., 3] + conv_state_out[..., 3] = x_t + return y, conv_state_out def qwen35_conv1d_decode_update( @@ -107,6 +53,7 @@ def qwen35_conv1d_decode_update( weight: torch.Tensor, *, activation: str = "silu", + backend: str = "auto", ) -> tuple[torch.Tensor, torch.Tensor]: """Single-token depthwise causal conv1d update. @@ -145,26 +92,28 @@ def qwen35_conv1d_decode_update( x_t = x_t.contiguous() conv_state = conv_state.contiguous() weight_2d = weight_2d.contiguous() - y = torch.empty_like(x_t) - - B, C = x_t.shape - cache = _get_compiled_kernel(B, C, kernel_size) - if "compiled" not in cache: - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - compiled = cute.compile( - _run_qwen35_conv1d_decode, - from_dlpack(x_t, assumed_align=16), - from_dlpack(conv_state, assumed_align=16), - from_dlpack(weight_2d, assumed_align=16), - from_dlpack(y, assumed_align=16), - B=B, - C=C, - K=kernel_size, - stream=stream, - options="--enable-tvm-ffi", + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_conv1d_decode") + and x_t.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_conv1d_decode is not available.") + + if use_cudac: + mixed_qkv_3d = x_t.unsqueeze(1).contiguous() + out_3d = torch.empty_like(mixed_qkv_3d) + conv_state_out = conv_state.clone() + cula_cuda.qwen35_conv1d_decode( + mixed_qkv_3d, + conv_state_out, + weight_2d, + out_3d, ) - cache["compiled"] = compiled + return out_3d.squeeze(1), conv_state_out - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - cache["compiled"](x_t, conv_state, weight_2d, y, stream) - return y, conv_state + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + return qwen35_conv1d_decode_reference(x_t, conv_state, weight_2d) diff --git a/cula/ops/qwen35_layout_decode.py b/cula/ops/qwen35_layout_decode.py new file mode 100644 index 00000000..1c63db41 --- /dev/null +++ b/cula/ops/qwen35_layout_decode.py @@ -0,0 +1,111 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3.5 layout decode wrapper.""" + +from __future__ import annotations + +import torch + +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig, infer_local_config + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def qwen35_layout_decode_reference( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + local_key_dim, _, local_num_k_heads = infer_local_config( + mixed_qkv_conv.shape[1], + local_num_v_heads, + config=config, + ) + + q_end = local_key_dim + k_end = q_end + local_key_dim + q_flat = mixed_qkv_conv[:, :q_end] + k_flat = mixed_qkv_conv[:, q_end:k_end] + v_flat = mixed_qkv_conv[:, k_end:] + + q = q_flat.view(tokens, local_num_k_heads, config.head_k_dim) + k = k_flat.view(tokens, local_num_k_heads, config.head_k_dim) + v = v_flat.view(tokens, local_num_v_heads, config.head_v_dim) + + repeat_factor = local_num_v_heads // local_num_k_heads + q_rep = q.repeat_interleave(repeat_factor, dim=1).contiguous() + k_rep = k.repeat_interleave(repeat_factor, dim=1).contiguous() + return q_rep, k_rep, v.contiguous(), a.contiguous(), b.contiguous() + + +def qwen35_layout_decode( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_layout_decode") + and mixed_qkv_conv.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_decode is not available.") + + if use_cudac: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + q_rep = torch.empty( + tokens, + local_num_v_heads, + config.head_k_dim, + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + k_rep = torch.empty_like(q_rep) + v = torch.empty( + tokens, + local_num_v_heads, + config.head_v_dim, + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + a_kernel = torch.empty_like(a) + b_kernel = torch.empty_like(b) + cula_cuda.qwen35_layout_decode( + mixed_qkv_conv.contiguous(), + a.contiguous(), + b.contiguous(), + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + ) + return q_rep, k_rep, v, a_kernel, b_kernel + + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + return qwen35_layout_decode_reference(mixed_qkv_conv, a, b, config=config) diff --git a/cula/ops/qwen35_scalar_kda_decode.py b/cula/ops/qwen35_scalar_kda_decode.py index db6ece66..41cff627 100644 --- a/cula/ops/qwen35_scalar_kda_decode.py +++ b/cula/ops/qwen35_scalar_kda_decode.py @@ -18,7 +18,10 @@ import torch -from cula.ops.kda_decode import kda_decode +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None def qwen35_scalar_kda_decode( @@ -32,6 +35,7 @@ def qwen35_scalar_kda_decode( recurrent_state: torch.Tensor, *, state_indices: torch.Tensor | None = None, + backend: str = "auto", ) -> tuple[torch.Tensor, torch.Tensor]: """Single-token scalar-gated delta-rule decode for Qwen3.5.""" if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: @@ -51,13 +55,50 @@ def qwen35_scalar_kda_decode( if A_log.shape != (HV,) or dt_bias.shape != (HV,): raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") - a_expanded = a.unsqueeze(-1).expand(N, 1, HV, K) - dt_bias_expanded = dt_bias[:, None].expand(HV, K).contiguous() state_indices = ( torch.arange(N, device=q.device, dtype=torch.int32) if state_indices is None else state_indices.to(device=q.device, dtype=torch.int32) ) + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_scalar_kda_decode") + and q.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_decode is not available.") + + if use_cudac: + q_rep = q.squeeze(1).contiguous() + k_rep = k.squeeze(1).contiguous() + v_rep = v.squeeze(1).contiguous() + a_kernel = a.squeeze(1).contiguous() + b_kernel = b.squeeze(1).contiguous() + out = torch.empty_like(v_rep) + recurrent_state_out = recurrent_state.clone() + cula_cuda.qwen35_scalar_kda_decode( + q_rep, + k_rep, + v_rep, + a_kernel, + b_kernel, + A_log.contiguous(), + dt_bias.contiguous(), + recurrent_state_out, + state_indices, + out, + ) + return out.unsqueeze(1), recurrent_state_out + + if backend not in ("auto", "generic_kda"): + raise ValueError(f"Unsupported backend={backend}") + + from cula.ops.kda_decode import kda_decode + + a_expanded = a.unsqueeze(-1).expand(N, 1, HV, K) + dt_bias_expanded = dt_bias[:, None].expand(HV, K).contiguous() o = kda_decode( A_log=A_log.contiguous(), dt_bias=dt_bias_expanded, diff --git a/cula/qwen35/__init__.py b/cula/qwen35/__init__.py index a8fd8351..4cb88c3d 100644 --- a/cula/qwen35/__init__.py +++ b/cula/qwen35/__init__.py @@ -15,13 +15,14 @@ """Qwen3.5-specific linear attention support built on top of cuLA primitives.""" from cula.qwen35.common import Qwen35LinearAttentionConfig -from cula.qwen35.runtime import ( - qwen35_linear_attention_decode, - qwen35_linear_attention_prefill, -) -__all__ = [ - "Qwen35LinearAttentionConfig", - "qwen35_linear_attention_prefill", - "qwen35_linear_attention_decode", -] +try: + from cula.qwen35.runtime import ( + qwen35_linear_attention_decode, + qwen35_linear_attention_prefill, + ) +except Exception: # pragma: no cover - optional runtime dependency during partial imports + qwen35_linear_attention_decode = None + qwen35_linear_attention_prefill = None + +__all__ = ["Qwen35LinearAttentionConfig", "qwen35_linear_attention_prefill", "qwen35_linear_attention_decode"] diff --git a/cula/qwen35/runtime.py b/cula/qwen35/runtime.py index 47f9548a..e633791e 100644 --- a/cula/qwen35/runtime.py +++ b/cula/qwen35/runtime.py @@ -16,9 +16,13 @@ from __future__ import annotations -import cuda.bindings.driver as cuda import torch +try: + import cuda.bindings.driver as cuda +except ImportError: # pragma: no cover - optional runtime dependency + cuda = None + from cula.qwen35.common import ( DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig, @@ -28,12 +32,15 @@ validate_state_tensors, ) from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_update +from cula.ops.qwen35_layout_decode import qwen35_layout_decode from cula.ops.qwen35_scalar_kda_decode import qwen35_scalar_kda_decode -_stream_cache: dict[tuple[str, int], cuda.CUstream] = {} +_stream_cache: dict[tuple[str, int], object] = {} -def _get_cached_stream(device: torch.device) -> cuda.CUstream: +def _get_cached_stream(device: torch.device) -> object: + if cuda is None: + raise RuntimeError("cuda.bindings.driver is not available in this environment.") stream_id = int(torch.cuda.current_stream(device=device).cuda_stream) cache_key = (str(device), stream_id) if cache_key not in _stream_cache: @@ -41,6 +48,99 @@ def _get_cached_stream(device: torch.device) -> cuda.CUstream: return _stream_cache[cache_key] +def _torch_qwen35_scalar_kda_decode_reference( + q_rep: torch.Tensor, + k_rep: torch.Tensor, + v: torch.Tensor, + a_kernel: torch.Tensor, + b_kernel: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pure torch reference for Qwen3.5 scalar-gated decode.""" + tokens, num_v_heads, head_k_dim = q_rep.shape + head_v_dim = v.shape[-1] + state_out = recurrent_state.clone() + out = torch.empty(tokens, num_v_heads, head_v_dim, device=q_rep.device, dtype=q_rep.dtype) + + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) + k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) + v_f = v.float() + a_f = a_kernel.float() + b_f = b_kernel.float() + + for token_idx in range(tokens): + pool_idx = int(state_indices[token_idx].item()) + for hv in range(num_v_heads): + state_kv = state_out[pool_idx, hv] + state_vk = state_kv.transpose(0, 1).contiguous() + + decay_pre = a_f[token_idx, hv] + dt_bias[hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(decay_pre)) + beta = torch.sigmoid(b_f[token_idx, hv]) + + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + + proj = state_vk @ k_vec + v_new = beta * (v_f[token_idx, hv] - proj) + state_vk_new = decay * state_vk + v_new.unsqueeze(1) * k_vec.unsqueeze(0) + out[token_idx, hv] = (state_vk_new @ q_vec).to(out.dtype) + state_out[pool_idx, hv] = state_vk_new.transpose(0, 1).contiguous() + + return out, state_out + + +def qwen35_linear_attention_decode_reference( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + conv_state: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure torch reference for the full Qwen3.5 decode chain.""" + tokens = mixed_qkv.shape[0] + if state_indices is None: + state_indices = torch.arange(tokens, device=mixed_qkv.device, dtype=torch.int32) + else: + state_indices = state_indices.to(device=mixed_qkv.device, dtype=torch.int32) + + conv_out, conv_state_out = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + backend="reference", + ) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode( + conv_out, + a, + b, + config=config, + backend="reference", + ) + core_attn_out, recurrent_state_out = _torch_qwen35_scalar_kda_decode_reference( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log.float(), + dt_bias.float(), + recurrent_state.float(), + state_indices, + ) + return core_attn_out.reshape(tokens, -1), conv_state_out, recurrent_state_out + + def qwen35_linear_attention_prefill( mixed_qkv: torch.Tensor, a: torch.Tensor, @@ -80,6 +180,7 @@ def qwen35_linear_attention_decode( conv_state: torch.Tensor, recurrent_state: torch.Tensor, state_indices: torch.Tensor | None = None, + backend: str = "auto", ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Qwen3.5 decode wrapper. @@ -100,7 +201,8 @@ def qwen35_linear_attention_decode( validate_mixed_qkv(mixed_qkv, config) validate_scalar_gate_inputs(a, b, config) validate_state_tensors(conv_state, recurrent_state, config) - _get_cached_stream(mixed_qkv.device) + if mixed_qkv.is_cuda: + _get_cached_stream(mixed_qkv.device) if mixed_qkv.shape[0] != a.shape[0]: raise ValueError(f"Token dimension mismatch, got mixed_qkv={tuple(mixed_qkv.shape)} a={tuple(a.shape)}") @@ -128,38 +230,52 @@ def qwen35_linear_attention_decode( if A_log.numel() != local_num_v_heads: raise ValueError(f"A_log must match local_num_v_heads={local_num_v_heads}, got {A_log.numel()}") + if backend == "auto" and not mixed_qkv.is_cuda: + backend = "reference" + + if backend == "reference": + return qwen35_linear_attention_decode_reference( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + ) + conv_out, conv_state_out = qwen35_conv1d_decode_update( mixed_qkv, conv_state, conv_weight, activation="silu", + backend=backend, ) - - q_end = local_key_dim - k_end = q_end + local_key_dim - q_flat = conv_out[:, :q_end] - k_flat = conv_out[:, q_end:k_end] - v_flat = conv_out[:, k_end:] - - q = q_flat.view(tokens, local_num_k_heads, config.head_k_dim) - k = k_flat.view(tokens, local_num_k_heads, config.head_k_dim) - v = v_flat.view(tokens, local_num_v_heads, config.head_v_dim) - - repeat_factor = local_num_v_heads // local_num_k_heads - q = q.repeat_interleave(repeat_factor, dim=1).unsqueeze(1).contiguous() - k = k.repeat_interleave(repeat_factor, dim=1).unsqueeze(1).contiguous() + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode( + conv_out, + a, + b, + config=config, + backend=backend, + ) + q = q_rep.unsqueeze(1).contiguous() + k = k_rep.unsqueeze(1).contiguous() v = v.unsqueeze(1).contiguous() core_attn_out, recurrent_state_out = qwen35_scalar_kda_decode( q=q, k=k, v=v, - a=a, - b=b, + a=a_kernel, + b=b_kernel, A_log=A_log, dt_bias=dt_bias, recurrent_state=recurrent_state, state_indices=state_indices, + backend=backend, ) core_attn_out = core_attn_out.reshape(tokens, local_value_dim) return core_attn_out, conv_state_out, recurrent_state_out diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py new file mode 100644 index 00000000..8ccfe39d --- /dev/null +++ b/tests/test_qwen35_decode.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_reference, qwen35_conv1d_decode_update +from cula.ops.qwen35_layout_decode import qwen35_layout_decode, qwen35_layout_decode_reference +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG +from cula.qwen35.runtime import qwen35_linear_attention_decode + + +def _device(): + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def make_inputs(tokens: int = 2, pool_size: int = 3, device: torch.device | None = None): + device = _device() if device is None else device + config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG + torch.manual_seed(0) + mixed_qkv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + conv_weight = torch.randn(config.conv_dim, config.conv_kernel_size, device=device, dtype=config.qkv_dtype) + conv_state = torch.randn(tokens, config.conv_dim, config.conv_kernel_size, device=device, dtype=config.qkv_dtype) + recurrent_state = torch.randn( + pool_size, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) % pool_size + return mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices + + +def manual_conv_decode(x_t: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor): + state_tail = conv_state[..., 1:].float() + window = torch.cat([state_tail, x_t.unsqueeze(-1).float()], dim=-1) + conv = (window * weight.float().unsqueeze(0)).sum(dim=-1) + y = torch.nn.functional.silu(conv).to(dtype=x_t.dtype) + state_new = conv_state.clone() + state_new[..., 0] = conv_state[..., 1] + state_new[..., 1] = conv_state[..., 2] + state_new[..., 2] = conv_state[..., 3] + state_new[..., 3] = x_t + return y, state_new + + +def manual_qwen35_decode_reference( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + conv_weight: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + conv_state: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, +): + config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG + conv_out, conv_state_out = manual_conv_decode(mixed_qkv, conv_state, conv_weight) + q_end = config.key_dim + k_end = q_end + config.key_dim + q = conv_out[:, :q_end].view(mixed_qkv.shape[0], config.num_k_heads, config.head_k_dim) + k = conv_out[:, q_end:k_end].view(mixed_qkv.shape[0], config.num_k_heads, config.head_k_dim) + v = conv_out[:, k_end:].view(mixed_qkv.shape[0], config.num_v_heads, config.head_v_dim) + q_rep = q.repeat_interleave(config.qk_repeat_factor, dim=1) + k_rep = k.repeat_interleave(config.qk_repeat_factor, dim=1) + + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) + k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) + v_f = v.float() + state_out = recurrent_state.clone() + out = torch.empty(mixed_qkv.shape[0], config.value_dim, device=mixed_qkv.device, dtype=mixed_qkv.dtype) + + for token_idx in range(mixed_qkv.shape[0]): + per_token = [] + pool_idx = int(state_indices[token_idx].item()) + for hv in range(config.num_v_heads): + state_kv = state_out[pool_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(a[token_idx, hv].float() + dt_bias[hv])) + beta = torch.sigmoid(b[token_idx, hv].float()) + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + proj = state_kv.transpose(0, 1) @ k_vec + v_new = beta * (v_f[token_idx, hv] - proj) + state_new_kv = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + per_token.append((state_new_kv.transpose(0, 1) @ q_vec).to(mixed_qkv.dtype)) + state_out[pool_idx, hv] = state_new_kv + out[token_idx] = torch.cat(per_token, dim=0) + return out, conv_state_out, state_out + + +@pytest.mark.parametrize("tokens", [1, 2]) +def test_qwen35_conv_decode_reference(tokens: int): + mixed_qkv, _, _, conv_weight, conv_state, _, _, _, _ = make_inputs(tokens=tokens) + y_ref, state_ref = manual_conv_decode(mixed_qkv, conv_state, conv_weight) + y_op, state_op = qwen35_conv1d_decode_update(mixed_qkv, conv_state, conv_weight, backend="reference") + assert torch.equal(y_ref, y_op) + assert torch.equal(state_ref, state_op) + y_ref2, state_ref2 = qwen35_conv1d_decode_reference(mixed_qkv, conv_state, conv_weight) + assert torch.equal(y_ref, y_ref2) + assert torch.equal(state_ref, state_ref2) + + +def test_qwen35_layout_decode_reference(): + mixed_qkv, a, b, _, _, _, _, _, _ = make_inputs(tokens=2) + q_rep_ref, k_rep_ref, v_ref, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv, a, b) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode(mixed_qkv, a, b, backend="reference") + assert torch.equal(q_rep_ref, q_rep) + assert torch.equal(k_rep_ref, k_rep) + assert torch.equal(v_ref, v) + assert torch.equal(a_ref, a_kernel) + assert torch.equal(b_ref, b_kernel) + + +@pytest.mark.parametrize("tokens", [1, 2]) +def test_qwen35_decode_reference_chain(tokens: int): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs(tokens=tokens) + out_ref, conv_state_ref, recurrent_state_ref = manual_qwen35_decode_reference( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state, + recurrent_state, + state_indices, + ) + out, conv_state_out, recurrent_state_out = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="reference", + ) + + assert torch.allclose(out_ref.float(), out.float(), atol=1e-5, rtol=1e-5) + assert torch.equal(conv_state_ref, conv_state_out) + assert torch.allclose(recurrent_state_ref, recurrent_state_out, atol=1e-6, rtol=1e-6) From de00a8b402972235b997f771f2e4530f6fdcacf2 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Fri, 5 Jun 2026 19:35:41 +0000 Subject: [PATCH 03/35] Fix Qwen3.5 decode extension build --- csrc/qwen35/decode/qwen35_layout_decode.cu | 2 ++ csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp | 8 ++++---- setup.py | 3 +++ 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_layout_decode.cu b/csrc/qwen35/decode/qwen35_layout_decode.cu index e601f17f..06d74931 100644 --- a/csrc/qwen35/decode/qwen35_layout_decode.cu +++ b/csrc/qwen35/decode/qwen35_layout_decode.cu @@ -22,6 +22,8 @@ #include #include +namespace { + void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); } diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp index 454dd146..dd10a34e 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp @@ -380,10 +380,10 @@ struct Qwen35ScalarKdaDecodeMainloop { // registers across both proj and update/out passes with acceptable // register pressure - const float a_val = static_cast(a_scalar()); - const float b_val = static_cast(b_scalar()); - const float A_log_val = static_cast(A_log_scalar()); - const float dt_bias_val = static_cast(dt_bias_scalar()); + const float a_val = static_cast(a_scalar); + const float b_val = static_cast(b_scalar); + const float A_log_val = static_cast(A_log_scalar); + const float dt_bias_val = static_cast(dt_bias_scalar); const float g = -expf(A_log_val) * softplusf_approx(a_val + dt_bias_val); const float decay = expf(g); diff --git a/setup.py b/setup.py index f7b11b95..411401ea 100644 --- a/setup.py +++ b/setup.py @@ -147,6 +147,9 @@ def get_nvcc_thread_args(): cuda_sources = [ "csrc/api/pybind.cu", + "csrc/qwen35/decode/qwen35_conv1d_decode.cu", + "csrc/qwen35/decode/qwen35_layout_decode.cu", + "csrc/qwen35/decode/qwen35_scalar_kda_decode.cu", ] if not DISABLE_SM100 or not DISABLE_SM103: cuda_sources.extend( From 8c3cb63c94bd3d7ae91df1be99cb8322ce9f1d2a Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sat, 6 Jun 2026 20:24:03 +0000 Subject: [PATCH 04/35] Align Qwen3.5 decode with FLA semantics --- .../decode/qwen35_scalar_kda_kernel.hpp | 5 +- .../decode/qwen35_scalar_kda_mainloop.hpp | 29 ++++++- cula/ops/qwen35_scalar_kda_decode.py | 17 ++++ cula/qwen35/runtime.py | 5 +- tests/test_qwen35_decode.py | 85 ++++++++++++++++++- 5 files changed, 132 insertions(+), 9 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 7628a3ce..85f7bd5d 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -51,9 +51,10 @@ struct Qwen35ScalarKdaDecodeKernel { // - q/k/v are staged once per CTA // - proj/out intermediates remain in fp32 // - recurrent state itself remains in fp32 global storage - alignas(16) scalar_t q_smem[kHeadDimQK]; - alignas(16) scalar_t k_smem[kHeadDimQK]; + alignas(16) float q_smem[kHeadDimQK]; + alignas(16) float k_smem[kHeadDimQK]; alignas(16) scalar_t v_smem[kHeadDimV]; + alignas(16) float norm_smem[2]; alignas(16) float proj_smem[kHeadDimV]; alignas(16) float out_smem[kHeadDimV]; }; diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp index dd10a34e..0c9401bc 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp @@ -392,13 +392,14 @@ struct Qwen35ScalarKdaDecodeMainloop { auto q_smem = make_tensor(make_smem_ptr(storage.q_smem), make_layout(make_shape(Int{}))); auto k_smem = make_tensor(make_smem_ptr(storage.k_smem), make_layout(make_shape(Int{}))); auto v_smem = make_tensor(make_smem_ptr(storage.v_smem), make_layout(make_shape(Int{}))); + auto norm_smem = make_tensor(make_smem_ptr(storage.norm_smem), make_layout(make_shape(Int<2>{}))); auto proj_smem = make_tensor(make_smem_ptr(storage.proj_smem), make_layout(make_shape(Int{}))); auto out_smem = make_tensor(make_smem_ptr(storage.out_smem), make_layout(make_shape(Int{}))); // Stage q/k/v once per CTA for the current decode token. for (int idx = tid; idx < kHeadDimQK; idx += num_threads) { - q_smem(idx) = q_vec(idx); - k_smem(idx) = k_vec(idx); + q_smem(idx) = static_cast(q_vec(idx)); + k_smem(idx) = static_cast(k_vec(idx)); } for (int idx = tid; idx < kHeadDimV; idx += num_threads) { v_smem(idx) = v_vec(idx); @@ -407,6 +408,27 @@ struct Qwen35ScalarKdaDecodeMainloop { } __syncthreads(); + if (tid == 0) { + float q_norm_sq = 0.f; + float k_norm_sq = 0.f; +#pragma unroll + for (int idx = 0; idx < kHeadDimQK; ++idx) { + const float q_val = q_smem(idx); + const float k_val = k_smem(idx); + q_norm_sq += q_val * q_val; + k_norm_sq += k_val * k_val; + } + norm_smem(0) = rsqrtf(q_norm_sq + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); + norm_smem(1) = rsqrtf(k_norm_sq + 1e-6f); + } + __syncthreads(); + + for (int idx = tid; idx < kHeadDimQK; idx += num_threads) { + q_smem(idx) = q_smem(idx) * norm_smem(0); + k_smem(idx) = k_smem(idx) * norm_smem(1); + } + __syncthreads(); + ThreadRowPlan row_plan = make_thread_row_plan(tid); // First concrete ownership model: @@ -427,7 +449,8 @@ struct Qwen35ScalarKdaDecodeMainloop { proj_smem(row_plan.v_row) = proj_row; const float v_val = static_cast(v_smem(row_plan.v_row)); - const float v_new_row = beta * (v_val - proj_row); + const float decayed_proj_row = decay * proj_row; + const float v_new_row = beta * (v_val - decayed_proj_row); float out_row = 0.f; for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { diff --git a/cula/ops/qwen35_scalar_kda_decode.py b/cula/ops/qwen35_scalar_kda_decode.py index 41cff627..a8b252ea 100644 --- a/cula/ops/qwen35_scalar_kda_decode.py +++ b/cula/ops/qwen35_scalar_kda_decode.py @@ -24,6 +24,22 @@ cula_cuda = None +def _validate_cudac_state_indices(state_indices: torch.Tensor, *, rows: int, pool_size: int) -> None: + if state_indices.ndim != 1 or state_indices.numel() != rows: + raise ValueError(f"state_indices must be 1D with {rows} entries, got {tuple(state_indices.shape)}") + if rows == 0: + return + min_idx = int(state_indices.min().item()) + max_idx = int(state_indices.max().item()) + if min_idx < 0 or max_idx >= pool_size: + raise ValueError(f"state_indices must be in [0, {pool_size}), got min={min_idx} max={max_idx}") + if torch.unique(state_indices).numel() != rows: + raise ValueError( + "backend='cudac' requires unique state_indices within one decode launch; " + "duplicate rows need a sequential decode path." + ) + + def qwen35_scalar_kda_decode( q: torch.Tensor, k: torch.Tensor, @@ -71,6 +87,7 @@ def qwen35_scalar_kda_decode( raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_decode is not available.") if use_cudac: + _validate_cudac_state_indices(state_indices, rows=N, pool_size=recurrent_state.shape[0]) q_rep = q.squeeze(1).contiguous() k_rep = k.squeeze(1).contiguous() v_rep = v.squeeze(1).contiguous() diff --git a/cula/qwen35/runtime.py b/cula/qwen35/runtime.py index e633791e..f518f0c8 100644 --- a/cula/qwen35/runtime.py +++ b/cula/qwen35/runtime.py @@ -65,7 +65,8 @@ def _torch_qwen35_scalar_kda_decode_reference( state_out = recurrent_state.clone() out = torch.empty(tokens, num_v_heads, head_v_dim, device=q_rep.device, dtype=q_rep.dtype) - q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) + scale = q_rep.shape[-1] ** -0.5 + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) * scale k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) v_f = v.float() a_f = a_kernel.float() @@ -84,7 +85,7 @@ def _torch_qwen35_scalar_kda_decode_reference( k_vec = k_f[token_idx, hv] q_vec = q_f[token_idx, hv] - proj = state_vk @ k_vec + proj = decay * (state_vk @ k_vec) v_new = beta * (v_f[token_idx, hv] - proj) state_vk_new = decay * state_vk + v_new.unsqueeze(1) * k_vec.unsqueeze(0) out[token_idx, hv] = (state_vk_new @ q_vec).to(out.dtype) diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py index 8ccfe39d..54dbd2f1 100644 --- a/tests/test_qwen35_decode.py +++ b/tests/test_qwen35_decode.py @@ -26,11 +26,26 @@ from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG from cula.qwen35.runtime import qwen35_linear_attention_decode +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + def _device(): return torch.device("cuda" if torch.cuda.is_available() else "cpu") +def _has_qwen35_cudac(): + return ( + torch.cuda.is_available() + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_conv1d_decode") + and hasattr(cula_cuda, "qwen35_layout_decode") + and hasattr(cula_cuda, "qwen35_scalar_kda_decode") + ) + + def make_inputs(tokens: int = 2, pool_size: int = 3, device: torch.device | None = None): device = _device() if device is None else device config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG @@ -88,7 +103,8 @@ def manual_qwen35_decode_reference( q_rep = q.repeat_interleave(config.qk_repeat_factor, dim=1) k_rep = k.repeat_interleave(config.qk_repeat_factor, dim=1) - q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) + scale = config.head_k_dim**-0.5 + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) * scale k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) v_f = v.float() state_out = recurrent_state.clone() @@ -103,7 +119,7 @@ def manual_qwen35_decode_reference( beta = torch.sigmoid(b[token_idx, hv].float()) k_vec = k_f[token_idx, hv] q_vec = q_f[token_idx, hv] - proj = state_kv.transpose(0, 1) @ k_vec + proj = decay * (state_kv.transpose(0, 1) @ k_vec) v_new = beta * (v_f[token_idx, hv] - proj) state_new_kv = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) per_token.append((state_new_kv.transpose(0, 1) @ q_vec).to(mixed_qkv.dtype)) @@ -165,3 +181,68 @@ def test_qwen35_decode_reference_chain(tokens: int): assert torch.allclose(out_ref.float(), out.float(), atol=1e-5, rtol=1e-5) assert torch.equal(conv_state_ref, conv_state_out) assert torch.allclose(recurrent_state_ref, recurrent_state_out, atol=1e-6, rtol=1e-6) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("tokens", [1, 2, 4]) +def test_qwen35_decode_cudac_matches_reference(tokens: int): + # Decode batches represent distinct active sequences, so keep state rows unique + # to avoid intentionally racing multiple token updates against one cache row. + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=tokens, + pool_size=max(tokens, 3), + device=torch.device("cuda"), + ) + out_ref, conv_state_ref, recurrent_state_ref = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="reference", + ) + out, conv_state_out, recurrent_state_out = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + assert torch.allclose(out_ref.float(), out.float(), atol=3e-2, rtol=3e-2) + assert torch.equal(conv_state_ref, conv_state_out) + assert torch.allclose(recurrent_state_ref, recurrent_state_out, atol=3e-5, rtol=3e-5) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +def test_qwen35_decode_cudac_rejects_duplicate_state_indices(): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, _ = make_inputs( + tokens=2, + pool_size=3, + device=torch.device("cuda"), + ) + state_indices = torch.zeros(2, device=mixed_qkv.device, dtype=torch.int32) + + with pytest.raises(ValueError, match="requires unique state_indices"): + qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) From 8c400835f76cc9257a79d2a7a671391df0fe1bcd Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Wed, 10 Jun 2026 03:38:17 +0000 Subject: [PATCH 05/35] Fuse Qwen3.5 layout into scalar KDA decode --- csrc/api/pybind.cu | 24 +++ csrc/qwen35/decode/qwen35_decode_common.cuh | 12 ++ .../qwen35/decode/qwen35_scalar_kda_decode.cu | 91 ++++++++++++ .../decode/qwen35_scalar_kda_kernel.hpp | 140 ++++++++++++++++++ cula/ops/qwen35_scalar_kda_decode.py | 91 ++++++++++++ cula/qwen35/runtime.py | 28 +++- tests/test_qwen35_decode.py | 49 ++++++ 7 files changed, 434 insertions(+), 1 deletion(-) diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index cd05e5c0..59bc3e71 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -135,6 +135,29 @@ qwen35_scalar_kda_decode( cula::qwen35::decode::run_qwen35_scalar_kda_decode(params); } +void +qwen35_layout_scalar_kda_decode( + at::Tensor mixed_qkv_conv, + at::Tensor a, + at::Tensor b, + at::Tensor A_log, + at::Tensor dt_bias, + at::Tensor recurrent_state, + at::Tensor pool_idx, + at::Tensor out) { + cula::qwen35::decode::LayoutScalarKdaDecodeParams params{ + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + }; + cula::qwen35::decode::run_qwen35_layout_scalar_kda_decode(params); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "cuLA"; #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) @@ -147,4 +170,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("qwen35_conv1d_decode", &qwen35_conv1d_decode); m.def("qwen35_layout_decode", &qwen35_layout_decode); m.def("qwen35_scalar_kda_decode", &qwen35_scalar_kda_decode); + m.def("qwen35_layout_scalar_kda_decode", &qwen35_layout_scalar_kda_decode); } diff --git a/csrc/qwen35/decode/qwen35_decode_common.cuh b/csrc/qwen35/decode/qwen35_decode_common.cuh index 53b64d65..8fad1378 100644 --- a/csrc/qwen35/decode/qwen35_decode_common.cuh +++ b/csrc/qwen35/decode/qwen35_decode_common.cuh @@ -66,8 +66,20 @@ struct ScalarKdaDecodeParams { at::Tensor out; // [N, 48, 128] }; +struct LayoutScalarKdaDecodeParams { + at::Tensor mixed_qkv_conv; // [N, 10240] + at::Tensor a; // [N, 48] + at::Tensor b; // [N, 48] + at::Tensor A_log; // [48], float32 + at::Tensor dt_bias; // [48], float32 + at::Tensor recurrent_state; // [pool, 48, 128, 128], float32 + at::Tensor pool_idx; // [N], int32 + at::Tensor out; // [N, 48, 128] +}; + void run_qwen35_conv1d_decode(ConvDecodeParams& params); void run_qwen35_layout_decode(LayoutDecodeParams& params); void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params); +void run_qwen35_layout_scalar_kda_decode(LayoutScalarKdaDecodeParams& params); } // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu index abed9854..a69f5bc6 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu +++ b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu @@ -131,4 +131,95 @@ void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params) { C10_CUDA_KERNEL_LAUNCH_CHECK(); } +void run_qwen35_layout_scalar_kda_decode(LayoutScalarKdaDecodeParams& params) { + const at::Tensor& mixed_qkv_conv = params.mixed_qkv_conv; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& A_log = params.A_log; + const at::Tensor& dt_bias = params.dt_bias; + const at::Tensor& recurrent_state = params.recurrent_state; + const at::Tensor& pool_idx = params.pool_idx; + const at::Tensor& out = params.out; + + TORCH_CHECK(mixed_qkv_conv.is_cuda(), "mixed_qkv_conv must be a CUDA tensor."); + const at::Device device = mixed_qkv_conv.device(); + + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(A_log, "A_log", device); + check_tensor_device(dt_bias, "dt_bias", device); + check_tensor_device(recurrent_state, "recurrent_state", device); + check_tensor_device(pool_idx, "pool_idx", device); + check_tensor_device(out, "out", device); + + TORCH_CHECK(mixed_qkv_conv.is_contiguous(), "mixed_qkv_conv must be contiguous."); + TORCH_CHECK(a.is_contiguous(), "a must be contiguous."); + TORCH_CHECK(b.is_contiguous(), "b must be contiguous."); + TORCH_CHECK(A_log.is_contiguous(), "A_log must be contiguous."); + TORCH_CHECK(dt_bias.is_contiguous(), "dt_bias must be contiguous."); + TORCH_CHECK(recurrent_state.is_contiguous(), "recurrent_state must be contiguous."); + TORCH_CHECK(pool_idx.is_contiguous(), "pool_idx must be contiguous."); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous."); + + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == a.scalar_type() && + mixed_qkv_conv.scalar_type() == b.scalar_type() && + mixed_qkv_conv.scalar_type() == out.scalar_type(), + "mixed_qkv_conv/a/b/out must share the same dtype."); + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == at::kHalf || mixed_qkv_conv.scalar_type() == at::kBFloat16, + "mixed_qkv_conv must be float16 or bfloat16."); + TORCH_CHECK(A_log.scalar_type() == at::kFloat, "A_log must be float32."); + TORCH_CHECK(dt_bias.scalar_type() == at::kFloat, "dt_bias must be float32."); + TORCH_CHECK(recurrent_state.scalar_type() == at::kFloat, "recurrent_state must be float32."); + TORCH_CHECK(pool_idx.scalar_type() == at::kInt, "pool_idx must be int32."); + + const int64_t token_count = mixed_qkv_conv.size(0); + TORCH_CHECK( + mixed_qkv_conv.dim() == 2 && + mixed_qkv_conv.sizes() == at::IntArrayRef({token_count, kMixedQKVDim}), + "mixed_qkv_conv must have shape [N, 10240]."); + TORCH_CHECK( + a.dim() == 2 && a.sizes() == at::IntArrayRef({token_count, kNumVHeads}), + "a must have shape [N, 48]."); + TORCH_CHECK( + b.dim() == 2 && b.sizes() == at::IntArrayRef({token_count, kNumVHeads}), + "b must have shape [N, 48]."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == kNumVHeads, "A_log must have shape [48]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == kNumVHeads, "dt_bias must have shape [48]."); + TORCH_CHECK( + recurrent_state.dim() == 4 && + recurrent_state.size(1) == kNumVHeads && + recurrent_state.size(2) == kHeadDimQK && + recurrent_state.size(3) == kHeadDimV, + "recurrent_state must have shape [pool, 48, 128, 128]."); + TORCH_CHECK(pool_idx.dim() == 1 && pool_idx.size(0) == token_count, "pool_idx must have shape [N]."); + TORCH_CHECK( + out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimV}), + "out must have shape [N, 48, 128]."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + mixed_qkv_conv.scalar_type(), + "launch_qwen35_layout_scalar_kda_decode_kernel", + [&] { + kernel::launch_qwen35_layout_scalar_kda_decode_kernel( + stream, + mixed_qkv_conv.data_ptr(), + a.data_ptr(), + b.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + recurrent_state.data_ptr(), + pool_idx.data_ptr(), + out.data_ptr(), + static_cast(token_count)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + } // namespace cula::qwen35::decode diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 85f7bd5d..1a767524 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -184,6 +184,95 @@ struct Qwen35ScalarKdaDecodeKernel { tid, kThreads); } + + template + CUTE_DEVICE static void run_layout_device( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count, + SharedStorage& storage) { + constexpr int kRepeatFactor = kNumVHeads / kNumQKHeads; + + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + if (token_idx >= token_count || hv >= kNumVHeads) { + return; + } + + const int mapped_h = hv / kRepeatFactor; + + auto qk_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kMixedQKVDim, kHeadDimQK, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kMixedQKVDim, kHeadDimV, Int<1>{})); + auto out_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kNumVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kNumVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_kv = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + const scalar_t* q_src = mixed_qkv_conv; + const scalar_t* k_src = mixed_qkv_conv + kQDim; + const scalar_t* v_src = mixed_qkv_conv + kQDim + kKDim; + + auto gQ = make_tensor(make_gmem_ptr(q_src), qk_src_layout); + auto gK = make_tensor(make_gmem_ptr(k_src), qk_src_layout); + auto gV = make_tensor(make_gmem_ptr(v_src), v_src_layout); + auto gO = make_tensor(make_gmem_ptr(out), out_layout); + auto gA = make_tensor(make_gmem_ptr(a), head_layout); + auto gB = make_tensor(make_gmem_ptr(b), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_kv = make_tensor(make_gmem_ptr(recurrent_state), state_layout_kv); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + (void)gH_kv; + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, mapped_h, _); + auto k_vec = gK(token_idx, mapped_h, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto a_scalar = gA(token_idx, hv); + auto b_scalar = gB(token_idx, hv); + auto A_log_scalar = gAlog(hv); + auto dt_bias_scalar = gDt(hv); + auto state_vk = gH_vk(state_row, hv, _, _); + + Mainloop::run( + q_vec, + k_vec, + v_vec, + a_scalar, + b_scalar, + A_log_scalar, + dt_bias_scalar, + state_vk, + out_vec, + storage, + tid, + kThreads); + } }; template > @@ -245,4 +334,55 @@ void launch_qwen35_scalar_kda_decode_kernel( token_count); } +template > +__global__ void qwen35_layout_scalar_kda_decode_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; + Qwen35ScalarKdaDecodeKernel::template run_layout_device( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count, + storage); +} + +template > +void launch_qwen35_layout_scalar_kda_decode_kernel( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_layout_scalar_kda_decode_kernel<<>>( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + } // namespace cula::qwen35::decode::kernel diff --git a/cula/ops/qwen35_scalar_kda_decode.py b/cula/ops/qwen35_scalar_kda_decode.py index a8b252ea..c32c30ee 100644 --- a/cula/ops/qwen35_scalar_kda_decode.py +++ b/cula/ops/qwen35_scalar_kda_decode.py @@ -24,6 +24,10 @@ cula_cuda = None +def has_qwen35_layout_scalar_kda_decode_cudac() -> bool: + return cula_cuda is not None and hasattr(cula_cuda, "qwen35_layout_scalar_kda_decode") + + def _validate_cudac_state_indices(state_indices: torch.Tensor, *, rows: int, pool_size: int) -> None: if state_indices.ndim != 1 or state_indices.numel() != rows: raise ValueError(f"state_indices must be 1D with {rows} entries, got {tuple(state_indices.shape)}") @@ -131,3 +135,90 @@ def qwen35_scalar_kda_decode( state_layout="kv", ) return o, recurrent_state + + +def qwen35_layout_scalar_kda_decode( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + *, + state_indices: torch.Tensor | None = None, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused Qwen3.5 layout decode + scalar-gated KDA decode.""" + if mixed_qkv_conv.ndim != 2: + raise ValueError(f"mixed_qkv_conv must be 2D, got {tuple(mixed_qkv_conv.shape)}") + if a.ndim == 3: + if a.shape[1] != 1: + raise ValueError(f"a sequence dim must be 1 for decode, got {tuple(a.shape)}") + a = a.squeeze(1) + if b.ndim == 3: + if b.shape[1] != 1: + raise ValueError(f"b sequence dim must be 1 for decode, got {tuple(b.shape)}") + b = b.squeeze(1) + + N = mixed_qkv_conv.shape[0] + if a.ndim != 2 or b.ndim != 2 or a.shape != b.shape or a.shape[0] != N: + raise ValueError(f"a/b must be [N, HV], got a={tuple(a.shape)} b={tuple(b.shape)}") + HV = a.shape[1] + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + + state_indices = ( + torch.arange(N, device=mixed_qkv_conv.device, dtype=torch.int32) + if state_indices is None + else state_indices.to(device=mixed_qkv_conv.device, dtype=torch.int32) + ) + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_layout_scalar_kda_decode") + and mixed_qkv_conv.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_scalar_kda_decode is not available.") + + if use_cudac: + _validate_cudac_state_indices(state_indices, rows=N, pool_size=recurrent_state.shape[0]) + out = torch.empty( + N, + HV, + recurrent_state.shape[-1], + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + recurrent_state_out = recurrent_state.clone() + cula_cuda.qwen35_layout_scalar_kda_decode( + mixed_qkv_conv.contiguous(), + a.contiguous(), + b.contiguous(), + A_log.contiguous(), + dt_bias.contiguous(), + recurrent_state_out, + state_indices, + out, + ) + return out.unsqueeze(1), recurrent_state_out + + if backend not in ("auto", "generic_kda"): + raise ValueError(f"Unsupported backend={backend}") + + from cula.ops.qwen35_layout_decode import qwen35_layout_decode_reference + + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode_reference(mixed_qkv_conv, a, b) + return qwen35_scalar_kda_decode( + q=q_rep.unsqueeze(1).contiguous(), + k=k_rep.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a_kernel, + b=b_kernel, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="generic_kda", + ) diff --git a/cula/qwen35/runtime.py b/cula/qwen35/runtime.py index f518f0c8..e2b5323e 100644 --- a/cula/qwen35/runtime.py +++ b/cula/qwen35/runtime.py @@ -33,7 +33,11 @@ ) from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_update from cula.ops.qwen35_layout_decode import qwen35_layout_decode -from cula.ops.qwen35_scalar_kda_decode import qwen35_scalar_kda_decode +from cula.ops.qwen35_scalar_kda_decode import ( + has_qwen35_layout_scalar_kda_decode_cudac, + qwen35_layout_scalar_kda_decode, + qwen35_scalar_kda_decode, +) _stream_cache: dict[tuple[str, int], object] = {} @@ -255,6 +259,28 @@ def qwen35_linear_attention_decode( activation="silu", backend=backend, ) + use_fused_layout_kda = ( + backend in ("auto", "cudac") + and mixed_qkv.is_cuda + and has_qwen35_layout_scalar_kda_decode_cudac() + ) + if backend == "cudac" and not use_fused_layout_kda: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_scalar_kda_decode is not available.") + + if use_fused_layout_kda: + core_attn_out, recurrent_state_out = qwen35_layout_scalar_kda_decode( + mixed_qkv_conv=conv_out, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend=backend, + ) + core_attn_out = core_attn_out.reshape(tokens, local_value_dim) + return core_attn_out, conv_state_out, recurrent_state_out + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode( conv_out, a, diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py index 54dbd2f1..d4364546 100644 --- a/tests/test_qwen35_decode.py +++ b/tests/test_qwen35_decode.py @@ -23,6 +23,7 @@ from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_reference, qwen35_conv1d_decode_update from cula.ops.qwen35_layout_decode import qwen35_layout_decode, qwen35_layout_decode_reference +from cula.ops.qwen35_scalar_kda_decode import qwen35_layout_scalar_kda_decode, qwen35_scalar_kda_decode from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG from cula.qwen35.runtime import qwen35_linear_attention_decode @@ -46,6 +47,10 @@ def _has_qwen35_cudac(): ) +def _has_qwen35_fused_layout_kda_cudac(): + return _has_qwen35_cudac() and hasattr(cula_cuda, "qwen35_layout_scalar_kda_decode") + + def make_inputs(tokens: int = 2, pool_size: int = 3, device: torch.device | None = None): device = _device() if device is None else device config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG @@ -224,6 +229,50 @@ def test_qwen35_decode_cudac_matches_reference(tokens: int): assert torch.allclose(recurrent_state_ref, recurrent_state_out, atol=3e-5, rtol=3e-5) +@pytest.mark.skipif(not _has_qwen35_fused_layout_kda_cudac(), reason="Qwen3.5 fused layout+KDA CUDA backend is not available") +@pytest.mark.parametrize("tokens", [1, 2, 4]) +def test_qwen35_fused_layout_kda_cudac_matches_unfused(tokens: int): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=tokens, + pool_size=max(tokens, 3), + device=torch.device("cuda"), + ) + conv_out, _ = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + backend="cudac", + ) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode(conv_out, a, b, backend="cudac") + out_unfused, state_unfused = qwen35_scalar_kda_decode( + q=q_rep.unsqueeze(1), + k=k_rep.unsqueeze(1), + v=v.unsqueeze(1), + a=a_kernel, + b=b_kernel, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + out_fused, state_fused = qwen35_layout_scalar_kda_decode( + mixed_qkv_conv=conv_out, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + assert torch.equal(out_unfused, out_fused) + assert torch.equal(state_unfused, state_fused) + + @pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") def test_qwen35_decode_cudac_rejects_duplicate_state_indices(): mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, _ = make_inputs( From c475de5070c190c25ed2c24ddd808ed5304b9a3d Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Wed, 10 Jun 2026 08:16:07 +0000 Subject: [PATCH 06/35] test qwen35 fused layout kda decode --- tests/test_qwen35_decode.py | 83 ++++++++++++++++++++++++++++++++++++- 1 file changed, 81 insertions(+), 2 deletions(-) diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py index d4364546..1dbe50d4 100644 --- a/tests/test_qwen35_decode.py +++ b/tests/test_qwen35_decode.py @@ -21,12 +21,17 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_reference, qwen35_conv1d_decode_update from cula.ops.qwen35_layout_decode import qwen35_layout_decode, qwen35_layout_decode_reference from cula.ops.qwen35_scalar_kda_decode import qwen35_layout_scalar_kda_decode, qwen35_scalar_kda_decode +from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_reference, qwen35_conv1d_decode_update from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG from cula.qwen35.runtime import qwen35_linear_attention_decode +try: + from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as triton_fused_sigmoid_update +except ImportError: + triton_fused_sigmoid_update = None + try: import cula.cudac as cula_cuda except ImportError: @@ -133,6 +138,47 @@ def manual_qwen35_decode_reference( return out, conv_state_out, state_out +def manual_qwen35_layout_scalar_kda_reference( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, +): + config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG + q_rep, k_rep, v, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv_conv, a, b) + + scale = config.head_k_dim**-0.5 + q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) * scale + k_f = torch.nn.functional.normalize(k_rep.float(), dim=-1) + v_f = v.float() + state_out = recurrent_state.clone() + out = torch.empty( + mixed_qkv_conv.shape[0], + q_rep.shape[1], + config.head_v_dim, + device=mixed_qkv_conv.device, + dtype=mixed_qkv_conv.dtype, + ) + + for token_idx in range(mixed_qkv_conv.shape[0]): + pool_idx = int(state_indices[token_idx].item()) + for hv in range(q_rep.shape[1]): + state_kv = state_out[pool_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(a_ref[token_idx, hv].float() + dt_bias[hv])) + beta = torch.sigmoid(b_ref[token_idx, hv].float()) + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v_f[token_idx, hv] - proj) + state_new_kv = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[token_idx, hv] = (state_new_kv.transpose(0, 1) @ q_vec).to(mixed_qkv_conv.dtype) + state_out[pool_idx, hv] = state_new_kv + return out.unsqueeze(1), state_out + + @pytest.mark.parametrize("tokens", [1, 2]) def test_qwen35_conv_decode_reference(tokens: int): mixed_qkv, _, _, conv_weight, conv_state, _, _, _, _ = make_inputs(tokens=tokens) @@ -231,7 +277,7 @@ def test_qwen35_decode_cudac_matches_reference(tokens: int): @pytest.mark.skipif(not _has_qwen35_fused_layout_kda_cudac(), reason="Qwen3.5 fused layout+KDA CUDA backend is not available") @pytest.mark.parametrize("tokens", [1, 2, 4]) -def test_qwen35_fused_layout_kda_cudac_matches_unfused(tokens: int): +def test_qwen35_fused_layout_kda_cudac_matches_reference_unfused_and_triton(tokens: int): mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( tokens=tokens, pool_size=max(tokens, 3), @@ -245,6 +291,15 @@ def test_qwen35_fused_layout_kda_cudac_matches_unfused(tokens: int): backend="cudac", ) q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode(conv_out, a, b, backend="cudac") + out_ref, state_ref = manual_qwen35_layout_scalar_kda_reference( + conv_out, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices, + ) out_unfused, state_unfused = qwen35_scalar_kda_decode( q=q_rep.unsqueeze(1), k=k_rep.unsqueeze(1), @@ -257,6 +312,25 @@ def test_qwen35_fused_layout_kda_cudac_matches_unfused(tokens: int): state_indices=state_indices, backend="cudac", ) + if triton_fused_sigmoid_update is not None: + state_triton = recurrent_state.clone() + out_triton = triton_fused_sigmoid_update( + A_log=A_log, + a=a_kernel.unsqueeze(1).contiguous(), + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_rep.unsqueeze(1).contiguous(), + k=k_rep.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + b=b_kernel.unsqueeze(1).contiguous(), + initial_state_source=state_triton, + initial_state_indices=state_indices, + scale=DEFAULT_QWEN35_LINEAR_ATTN_CONFIG.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + is_kda=False, + ) out_fused, state_fused = qwen35_layout_scalar_kda_decode( mixed_qkv_conv=conv_out, a=a, @@ -269,8 +343,13 @@ def test_qwen35_fused_layout_kda_cudac_matches_unfused(tokens: int): ) torch.cuda.synchronize() + assert torch.allclose(out_ref.float(), out_fused.float(), atol=3e-2, rtol=3e-2) + assert torch.allclose(state_ref, state_fused, atol=3e-5, rtol=3e-5) assert torch.equal(out_unfused, out_fused) assert torch.equal(state_unfused, state_fused) + if triton_fused_sigmoid_update is not None: + assert torch.allclose(out_triton.float(), out_fused.float(), atol=3e-2, rtol=3e-2) + assert torch.allclose(state_triton, state_fused, atol=3e-5, rtol=3e-5) @pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") From 5d1b495e64865c535e85c8e77a8f2525996c00f9 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Wed, 10 Jun 2026 08:46:44 +0000 Subject: [PATCH 07/35] benchmark qwen35 decode paths --- benchmarks/bench_qwen35_decode.py | 633 ++++++++++++++++++++++++++++++ 1 file changed, 633 insertions(+) create mode 100755 benchmarks/bench_qwen35_decode.py diff --git a/benchmarks/bench_qwen35_decode.py b/benchmarks/bench_qwen35_decode.py new file mode 100755 index 00000000..932f757b --- /dev/null +++ b/benchmarks/bench_qwen35_decode.py @@ -0,0 +1,633 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark Qwen3.5 decode on the active CUDA device. + +Two timing scopes are reported: + - native_core: direct native scalar GDN decode op. + - triton_core: FLA/SGLang-style fused_sigmoid_gating_delta_rule_update + Triton decode op vendored in cuLA. + - sglang_core: fused_sigmoid_gating_delta_rule_update from SGLang, when + available from the installed package or --sglang-path. + - fused_layout_kda: direct cuLA fused Qwen3.5 layout + scalar KDA decode op. + - sglang_packed: SGLang packed Qwen3.5 layout + recurrent update op, when + available from the installed package or --sglang-path. + - full: cuLA Python Qwen3.5 decode chain, including conv + layout + core. + +State buffers are reset before each timed iteration and the reset copy is not +included in the event timing window. +""" + +from __future__ import annotations + +import argparse +import csv +import importlib +import importlib.util +import inspect +import pathlib +import statistics +import sys +import time +from collections.abc import Callable + +import torch + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +import cula.cudac as cula_cuda +from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as triton_fused_sigmoid_update +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as CONFIG +from cula.qwen35.runtime import qwen35_linear_attention_decode + +SGLANG_CORE_MODULES = [ + "sglang.srt.layers.attention.linear.kernels.gdn_triton", + "sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent", +] +SGLANG_CORE_FILES = [ + pathlib.Path("sglang/srt/layers/attention/linear/kernels/gdn_triton.py"), + pathlib.Path("sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py"), +] +SGLANG_PACKED_MODULES = [ + "sglang.srt.layers.attention.fla.fused_recurrent", + "sglang.srt.layers.attention.linear.kernels.gdn_triton", +] +SGLANG_PACKED_FILES = [ + pathlib.Path("sglang/srt/layers/attention/fla/fused_recurrent.py"), + pathlib.Path("sglang/srt/layers/attention/linear/kernels/gdn_triton.py"), +] + + +def accelerator_device() -> torch.device: + if torch.cuda.is_available(): + return torch.device("cuda") + raise RuntimeError("No CUDA accelerator is available.") + + +def accelerator_name(device: torch.device) -> str: + if device.type != "cuda": + raise ValueError(f"Unsupported device={device}") + return torch.cuda.get_device_name(device.index or 0) + + +def synchronize(device: torch.device) -> None: + if device.type != "cuda": + raise ValueError(f"Unsupported device={device}") + torch.cuda.synchronize() + + +def benchmark_accel_fn( + fn: Callable[[], object], + *, + device: torch.device, + setup_fn: Callable[[], None] | None, + warmup: int, + rep: int, +) -> float: + for _ in range(warmup): + if setup_fn is not None: + setup_fn() + fn() + synchronize(device) + + times: list[float] = [] + try: + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + for i in range(rep): + if setup_fn is not None: + setup_fn() + starts[i].record() + fn() + ends[i].record() + synchronize(device) + times = [s.elapsed_time(e) for s, e in zip(starts, ends)] + except Exception: + for _ in range(rep): + if setup_fn is not None: + setup_fn() + synchronize(device) + t0 = time.perf_counter() + fn() + synchronize(device) + times.append((time.perf_counter() - t0) * 1000.0) + + if not times: + return 0.0 + if len(times) < 4: + return statistics.mean(times) + times = sorted(times) + iqr = times[len(times) // 4 : 3 * len(times) // 4] + return statistics.mean(iqr) + + +def make_full_inputs(tokens: int, device: torch.device, seed: int): + torch.manual_seed(seed) + pool_size = max(tokens, 1) + mixed_qkv = torch.randn(tokens, CONFIG.conv_dim, device=device, dtype=CONFIG.qkv_dtype) + a = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + b = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + conv_weight = torch.randn(CONFIG.conv_dim, CONFIG.conv_kernel_size, device=device, dtype=CONFIG.qkv_dtype) + conv_state = torch.randn( + tokens, + CONFIG.conv_dim, + CONFIG.conv_kernel_size, + device=device, + dtype=CONFIG.qkv_dtype, + ) + recurrent_state = torch.randn( + pool_size, + CONFIG.num_v_heads, + CONFIG.head_k_dim, + CONFIG.head_v_dim, + device=device, + dtype=CONFIG.state_dtype, + ) * 0.01 + A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + return mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices + + +def make_fused_layout_kda_inputs(tokens: int, device: torch.device, seed: int): + torch.manual_seed(seed) + mixed_qkv_conv = torch.randn(tokens, CONFIG.conv_dim, device=device, dtype=CONFIG.qkv_dtype) + a = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + b = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state = torch.randn( + tokens, + CONFIG.num_v_heads, + CONFIG.head_k_dim, + CONFIG.head_v_dim, + device=device, + dtype=CONFIG.state_dtype, + ) * 0.01 + state_work = torch.empty_like(state) + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + out = torch.empty(tokens, CONFIG.num_v_heads, CONFIG.head_v_dim, device=device, dtype=CONFIG.qkv_dtype) + return mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out + + +def make_core_inputs(tokens: int, device: torch.device, seed: int): + torch.manual_seed(seed) + q = torch.randn(tokens, CONFIG.num_v_heads, CONFIG.head_k_dim, device=device, dtype=CONFIG.qkv_dtype) + k = torch.randn(tokens, CONFIG.num_v_heads, CONFIG.head_k_dim, device=device, dtype=CONFIG.qkv_dtype) + v = torch.randn(tokens, CONFIG.num_v_heads, CONFIG.head_v_dim, device=device, dtype=CONFIG.qkv_dtype) + a = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + b = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state = torch.randn( + tokens, + CONFIG.num_v_heads, + CONFIG.head_k_dim, + CONFIG.head_v_dim, + device=device, + dtype=CONFIG.state_dtype, + ) * 0.01 + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + out = torch.empty_like(v) + state_work = torch.empty_like(state) + return q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out + + +def _add_sglang_import_roots(sglang_path: pathlib.Path | None) -> None: + if sglang_path is not None: + for import_root in (sglang_path, sglang_path / "python"): + if import_root.exists(): + sys.path.insert(0, str(import_root)) + + +def _resolve_sglang_symbol( + *, + sglang_path: pathlib.Path | None, + module_names: list[str], + file_paths: list[pathlib.Path], + symbol_names: list[str], +): + _add_sglang_import_roots(sglang_path) + + import_errors = [] + for module_name in module_names: + try: + module = importlib.import_module(module_name) + except (ImportError, PermissionError, ModuleNotFoundError) as exc: + import_errors.append(f"{module_name}: {type(exc).__name__}: {exc}") + continue + for symbol_name in symbol_names: + if hasattr(module, symbol_name): + return getattr(module, symbol_name), f"{module_name}.{symbol_name}" + + if sglang_path is not None: + candidates: list[pathlib.Path] = [] + for rel_path in file_paths: + candidates.extend([sglang_path / rel_path, sglang_path / "python" / rel_path]) + for idx, path in enumerate(candidates): + if path.exists(): + spec = importlib.util.spec_from_file_location(f"_sglang_qwen35_decode_provider_{idx}", path) + if spec is None or spec.loader is None: + continue + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + for symbol_name in symbol_names: + if hasattr(module, symbol_name): + return getattr(module, symbol_name), f"{path}:{symbol_name}" + raise RuntimeError( + f"Could not find any of {symbol_names} under --sglang-path={sglang_path}. " + "Pass the SGLang repo root or its python/ directory. " + f"Import errors: {'; '.join(import_errors) or 'none'}" + ) + + return None, None + + +def resolve_sglang_core_update(sglang_path: pathlib.Path | None): + """Return SGLang's scalar-gated recurrent update function if available.""" + return _resolve_sglang_symbol( + sglang_path=sglang_path, + module_names=SGLANG_CORE_MODULES, + file_paths=SGLANG_CORE_FILES, + symbol_names=["fused_sigmoid_gating_delta_rule_update"], + ) + + +def resolve_sglang_packed_decode(sglang_path: pathlib.Path | None): + """Return SGLang's packed layout + recurrent decode function if available.""" + return _resolve_sglang_symbol( + sglang_path=sglang_path, + module_names=SGLANG_PACKED_MODULES, + file_paths=SGLANG_PACKED_FILES, + symbol_names=[ + "fused_recurrent_gated_delta_rule_packed_decode", + "fused_recurrent_gated_delta_rule_packed_decode_cpu", + ], + ) + + +def call_with_supported_kwargs(fn: Callable, **kwargs): + """Call a provider while tolerating minor SGLang signature drift.""" + try: + signature = inspect.signature(fn) + except (TypeError, ValueError): + return fn(**kwargs) + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): + return fn(**kwargs) + filtered = {name: value for name, value in kwargs.items() if name in signature.parameters} + return fn(**filtered) + + +def bench_native_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: + q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_core_inputs(tokens, device, seed) + + def setup() -> None: + state_work.copy_(state) + + def run() -> None: + cula_cuda.qwen35_scalar_kda_decode( + q, + k, + v, + a, + b, + A_log, + dt_bias, + state_work, + state_indices, + out, + ) + + return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) + + +def bench_triton_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: + q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, _ = make_core_inputs(tokens, device, seed) + q_4d = q.unsqueeze(1).contiguous() + k_4d = k.unsqueeze(1).contiguous() + v_4d = v.unsqueeze(1).contiguous() + a_3d = a.unsqueeze(1).contiguous() + b_3d = b.unsqueeze(1).contiguous() + + def setup() -> None: + state_work.copy_(state) + + def run() -> None: + triton_fused_sigmoid_update( + A_log=A_log, + a=a_3d, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_4d, + k=k_4d, + v=v_4d, + b=b_3d, + initial_state_source=state_work, + initial_state_indices=state_indices, + scale=CONFIG.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + is_kda=False, + ) + + return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) + + +def bench_fused_layout_kda(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: + mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_fused_layout_kda_inputs( + tokens, device, seed + ) + + def setup() -> None: + state_work.copy_(state) + + def run() -> None: + cula_cuda.qwen35_layout_scalar_kda_decode( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + state_work, + state_indices, + out, + ) + + return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) + + +def bench_sglang_core( + tokens: int, + device: torch.device, + warmup: int, + rep: int, + seed: int, + sglang_fused_update: Callable, +) -> float: + q, k, v, a, b, A_log, dt_bias, state, _, state_indices, _ = make_core_inputs(tokens, device, seed) + q_4d = q.unsqueeze(1).contiguous() + k_4d = k.unsqueeze(1).contiguous() + v_4d = v.unsqueeze(1).contiguous() + a_3d = a.unsqueeze(1).contiguous() + b_3d = b.unsqueeze(1).contiguous() + state_vk = state.transpose(-1, -2).contiguous() + state_vk_work = torch.empty_like(state_vk) + + def setup() -> None: + state_vk_work.copy_(state_vk) + + def run() -> None: + call_with_supported_kwargs( + sglang_fused_update, + A_log=A_log, + a=a_3d, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_4d, + k=k_4d, + v=v_4d, + b=b_3d, + initial_state_source=state_vk_work, + initial_state_indices=state_indices, + scale=CONFIG.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + cu_seqlens=None, + is_kda=False, + ) + + return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) + + +def bench_sglang_packed_layout_kda( + tokens: int, + device: torch.device, + warmup: int, + rep: int, + seed: int, + sglang_packed_decode: Callable, +) -> float: + mixed_qkv_conv, a, b, A_log, dt_bias, state, _, state_indices, _ = make_fused_layout_kda_inputs(tokens, device, seed) + state_vk = state.transpose(-1, -2).contiguous() + state_vk_work = torch.empty_like(state_vk) + out = torch.empty(tokens, 1, CONFIG.num_v_heads, CONFIG.head_v_dim, device=device, dtype=CONFIG.qkv_dtype) + + def setup() -> None: + state_vk_work.copy_(state_vk) + + def run() -> None: + call_with_supported_kwargs( + sglang_packed_decode, + mixed_qkv=mixed_qkv_conv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + scale=CONFIG.head_k_dim**-0.5, + initial_state=state_vk_work, + out=out, + ssm_state_indices=state_indices, + use_qk_l2norm_in_kernel=True, + ) + + return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) + + +def bench_full(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: + inputs = make_full_inputs(tokens, device, seed) + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = inputs + conv_state_work = torch.empty_like(conv_state) + recurrent_state_work = torch.empty_like(recurrent_state) + + def setup() -> None: + conv_state_work.copy_(conv_state) + recurrent_state_work.copy_(recurrent_state) + + def run() -> None: + qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + conv_state=conv_state_work, + recurrent_state=recurrent_state_work, + state_indices=state_indices, + backend="cudac", + ) + + return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Benchmark cuLA Qwen3.5 decode.") + parser.add_argument("--tokens", nargs="+", type=int, default=[1, 2, 4, 8, 16, 32, 64, 128]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--scope", choices=["core", "fused", "full", "both"], default="both") + parser.add_argument("--skip-triton", action="store_true", help="Skip the vendored Triton core timing.") + parser.add_argument("--skip-sglang", action="store_true", help="Do not try the SGLang kernel provider.") + parser.add_argument("--require-sglang", action="store_true", help="Fail if the SGLang kernel provider is unavailable.") + parser.add_argument("--sglang-path", type=pathlib.Path, default=None, help="SGLang repo root or python/ directory.") + parser.add_argument("--csv", type=pathlib.Path, default=None) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + device = accelerator_device() + rows: list[dict[str, object]] = [] + sglang_fused_update = None + sglang_core_source = None + sglang_packed_decode = None + sglang_packed_source = None + if not args.skip_sglang: + sglang_fused_update, sglang_core_source = resolve_sglang_core_update(args.sglang_path) + sglang_packed_decode, sglang_packed_source = resolve_sglang_packed_decode(args.sglang_path) + if args.require_sglang and (sglang_fused_update is None or sglang_packed_decode is None): + raise RuntimeError("SGLang core and packed decode providers must both be available.") + + print(f"device={device} name={accelerator_name(device)} torch={torch.__version__}") + print(f"qwen35: HV={CONFIG.num_v_heads} K={CONFIG.head_k_dim} V={CONFIG.head_v_dim} conv_dim={CONFIG.conv_dim}") + print(f"sglang_core_provider={sglang_core_source or 'unavailable'}") + print(f"sglang_packed_provider={sglang_packed_source or 'unavailable'}") + print("| tokens | native_core_ms | triton_core_ms | sglang_core_ms | fused_layout_kda_ms | sglang_packed_ms | full_ms | triton/native | sglang/native | packed/fused | native_us_per_token | triton_us_per_token | sglang_us_per_token | fused_us_per_token | packed_us_per_token | full_us_per_token |") + print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|") + + for tokens in args.tokens: + native_core_ms = None + triton_core_ms = None + sglang_core_ms = None + fused_layout_kda_ms = None + sglang_packed_ms = None + full_ms = None + if args.scope in ("core", "both"): + native_core_ms = bench_native_core(tokens, device, args.warmup, args.rep, args.seed) + if not args.skip_triton: + triton_core_ms = bench_triton_core(tokens, device, args.warmup, args.rep, args.seed) + if sglang_fused_update is not None: + sglang_core_ms = bench_sglang_core( + tokens, + device, + args.warmup, + args.rep, + args.seed, + sglang_fused_update, + ) + if args.scope in ("fused", "both"): + fused_layout_kda_ms = bench_fused_layout_kda(tokens, device, args.warmup, args.rep, args.seed) + if sglang_packed_decode is not None: + sglang_packed_ms = bench_sglang_packed_layout_kda( + tokens, + device, + args.warmup, + args.rep, + args.seed, + sglang_packed_decode, + ) + if args.scope in ("full", "both"): + full_ms = bench_full(tokens, device, args.warmup, args.rep, args.seed) + + native_core_us = None if native_core_ms is None else native_core_ms * 1000.0 / tokens + triton_core_us = None if triton_core_ms is None else triton_core_ms * 1000.0 / tokens + sglang_core_us = None if sglang_core_ms is None else sglang_core_ms * 1000.0 / tokens + fused_layout_kda_us = None if fused_layout_kda_ms is None else fused_layout_kda_ms * 1000.0 / tokens + sglang_packed_us = None if sglang_packed_ms is None else sglang_packed_ms * 1000.0 / tokens + full_us = None if full_ms is None else full_ms * 1000.0 / tokens + triton_ratio = None + if native_core_ms is not None and triton_core_ms is not None and native_core_ms > 0: + triton_ratio = triton_core_ms / native_core_ms + sglang_ratio = None + if native_core_ms is not None and sglang_core_ms is not None and native_core_ms > 0: + sglang_ratio = sglang_core_ms / native_core_ms + packed_ratio = None + if fused_layout_kda_ms is not None and sglang_packed_ms is not None and fused_layout_kda_ms > 0: + packed_ratio = sglang_packed_ms / fused_layout_kda_ms + print( + f"| {tokens} | " + f"{'n/a' if native_core_ms is None else f'{native_core_ms:.4f}'} | " + f"{'n/a' if triton_core_ms is None else f'{triton_core_ms:.4f}'} | " + f"{'n/a' if sglang_core_ms is None else f'{sglang_core_ms:.4f}'} | " + f"{'n/a' if fused_layout_kda_ms is None else f'{fused_layout_kda_ms:.4f}'} | " + f"{'n/a' if sglang_packed_ms is None else f'{sglang_packed_ms:.4f}'} | " + f"{'n/a' if full_ms is None else f'{full_ms:.4f}'} | " + f"{'n/a' if triton_ratio is None else f'{triton_ratio:.2f}x'} | " + f"{'n/a' if sglang_ratio is None else f'{sglang_ratio:.2f}x'} | " + f"{'n/a' if packed_ratio is None else f'{packed_ratio:.2f}x'} | " + f"{'n/a' if native_core_us is None else f'{native_core_us:.2f}'} | " + f"{'n/a' if triton_core_us is None else f'{triton_core_us:.2f}'} | " + f"{'n/a' if sglang_core_us is None else f'{sglang_core_us:.2f}'} | " + f"{'n/a' if fused_layout_kda_us is None else f'{fused_layout_kda_us:.2f}'} | " + f"{'n/a' if sglang_packed_us is None else f'{sglang_packed_us:.2f}'} | " + f"{'n/a' if full_us is None else f'{full_us:.2f}'} |" + ) + rows.append( + { + "tokens": tokens, + "native_core_ms": native_core_ms, + "triton_core_ms": triton_core_ms, + "sglang_core_ms": sglang_core_ms, + "fused_layout_kda_ms": fused_layout_kda_ms, + "sglang_packed_ms": sglang_packed_ms, + "full_ms": full_ms, + "triton_over_native": triton_ratio, + "sglang_over_native": sglang_ratio, + "sglang_packed_over_fused": packed_ratio, + "native_core_us_per_token": native_core_us, + "triton_core_us_per_token": triton_core_us, + "sglang_core_us_per_token": sglang_core_us, + "fused_layout_kda_us_per_token": fused_layout_kda_us, + "sglang_packed_us_per_token": sglang_packed_us, + "full_us_per_token": full_us, + } + ) + + if args.csv is not None: + args.csv.parent.mkdir(parents=True, exist_ok=True) + with args.csv.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "tokens", + "native_core_ms", + "triton_core_ms", + "sglang_core_ms", + "fused_layout_kda_ms", + "sglang_packed_ms", + "full_ms", + "triton_over_native", + "sglang_over_native", + "sglang_packed_over_fused", + "native_core_us_per_token", + "triton_core_us_per_token", + "sglang_core_us_per_token", + "fused_layout_kda_us_per_token", + "sglang_packed_us_per_token", + "full_us_per_token", + ], + ) + writer.writeheader() + writer.writerows(rows) + print(f"wrote {args.csv}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 0a2579b7bafad247a8af1d35a41be0af03769bd0 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Wed, 10 Jun 2026 13:04:13 +0000 Subject: [PATCH 08/35] Add Qwen3.5 prefill CUDA path --- benchmarks/bench_qwen35_prefill.py | 290 ++++++++++++++++++ csrc/api/pybind.cu | 55 ++++ csrc/qwen35/prefill/qwen35_layout_prefill.cu | 127 ++++++++ .../prefill/qwen35_layout_prefill_kernel.hpp | 141 +++++++++ csrc/qwen35/prefill/qwen35_prefill_common.cuh | 60 ++++ .../prefill/qwen35_scalar_kda_prefill.cu | 171 +++++++++++ .../qwen35_scalar_kda_prefill_kernel.hpp | 273 +++++++++++++++++ cula/ops/qwen35_conv1d_prefill.py | 73 ++++- cula/ops/qwen35_layout_prefill.py | 97 ++++++ cula/ops/qwen35_scalar_kda_prefill.py | 112 ++++++- cula/qwen35/runtime.py | 71 ++++- setup.py | 2 + tests/test_qwen35_prefill.py | 175 +++++++++++ 13 files changed, 1633 insertions(+), 14 deletions(-) create mode 100644 benchmarks/bench_qwen35_prefill.py create mode 100644 csrc/qwen35/prefill/qwen35_layout_prefill.cu create mode 100644 csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp create mode 100644 csrc/qwen35/prefill/qwen35_prefill_common.cuh create mode 100644 csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu create mode 100644 csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp create mode 100644 cula/ops/qwen35_layout_prefill.py create mode 100644 tests/test_qwen35_prefill.py diff --git a/benchmarks/bench_qwen35_prefill.py b/benchmarks/bench_qwen35_prefill.py new file mode 100644 index 00000000..5e7dbe7b --- /dev/null +++ b/benchmarks/bench_qwen35_prefill.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark Qwen3.5 prefill kernels. + +Reports: + - layout: cuLA Qwen3.5 prefill layout split/repeat kernel + - scalar_kda: cuLA Qwen3.5 scalar-gated KDA prefill kernel + - fla_gdr: optional FLA chunk_gated_delta_rule baseline + - sgl_gdr: optional SGLang vendored Triton chunk_gated_delta_rule baseline + +Baselines are optional. SGLang Qwen3.5 prefill uses the same chunked gated +delta rule family in its Triton GDN kernel; decode uses a recurrent packed +kernel instead. +""" + +from __future__ import annotations + +import argparse +import importlib +import inspect +import pathlib +import statistics +import sys +from collections.abc import Callable + +import torch + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as CONFIG + + +def benchmark_cuda_fn(fn: Callable[[], object], *, warmup: int, rep: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + for idx in range(rep): + starts[idx].record() + fn() + ends[idx].record() + torch.cuda.synchronize() + + times = [start.elapsed_time(end) for start, end in zip(starts, ends)] + if len(times) <= 2: + return statistics.mean(times) + times = sorted(times) + return statistics.mean(times[len(times) // 4 : 3 * len(times) // 4]) + + +def error_stats(ref: torch.Tensor, out: torch.Tensor) -> tuple[float, float, float]: + ref_f = ref.float() + out_f = out.float() + diff = (ref_f - out_f).abs() + rmse = diff.square().mean().sqrt().item() + ref_rms = ref_f.square().mean().sqrt().item() + rel_rms = rmse / (ref_rms + 1.0e-8) + rel_max = diff.max().item() / (ref_f.abs().max().item() + 1.0e-8) + mean_abs = diff.mean().item() + return rel_rms, rel_max, mean_abs + + +def resolve_fla_chunk_gdr(): + try: + module = importlib.import_module("fla.ops.gated_delta_rule") + except ImportError as exc: + return None, f"cannot import fla.ops.gated_delta_rule: {exc}" + if not hasattr(module, "chunk_gated_delta_rule"): + return None, "fla.ops.gated_delta_rule has no chunk_gated_delta_rule" + return module.chunk_gated_delta_rule, "fla.ops.gated_delta_rule.chunk_gated_delta_rule" + + +def resolve_sgl_chunk_gdr(sglang_path: pathlib.Path | None): + if sglang_path is not None: + for root in (sglang_path, sglang_path / "python"): + if root.exists(): + sys.path.insert(0, str(root)) + try: + module = importlib.import_module("sglang.srt.layers.attention.fla.chunk") + except ImportError as exc: + return None, f"cannot import sglang.srt.layers.attention.fla.chunk: {exc}" + if not hasattr(module, "chunk_gated_delta_rule"): + return None, "sglang.srt.layers.attention.fla.chunk has no chunk_gated_delta_rule" + return module.chunk_gated_delta_rule, "sglang.srt.layers.attention.fla.chunk.chunk_gated_delta_rule" + + +def make_inputs(batch: int, seq_len: int, *, device: torch.device, seed: int): + torch.manual_seed(seed) + q = torch.randn(batch, seq_len, CONFIG.num_v_heads, CONFIG.head_k_dim, device=device, dtype=CONFIG.qkv_dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(batch, seq_len, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + b = torch.randn(batch, seq_len, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) + beta = torch.sigmoid(b.float()).to(dtype=CONFIG.qkv_dtype) + A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + log_gate = (-torch.exp(A_log).view(1, 1, -1) * torch.nn.functional.softplus(a.float() + dt_bias.view(1, 1, -1))).to( + dtype=CONFIG.qkv_dtype + ) + initial_state = torch.randn( + batch, + CONFIG.num_v_heads, + CONFIG.head_k_dim, + CONFIG.head_v_dim, + device=device, + dtype=torch.float32, + ) * 0.01 + mixed_qkv_conv = torch.randn(batch * seq_len, CONFIG.conv_dim, device=device, dtype=CONFIG.qkv_dtype) + a_flat = a.reshape(batch * seq_len, CONFIG.num_v_heads).contiguous() + b_flat = b.reshape(batch * seq_len, CONFIG.num_v_heads).contiguous() + return q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, mixed_qkv_conv, a_flat, b_flat + + +def run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state): + return qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="cudac", + ) + + +def run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices): + # SGLang/FLA GDR chunk kernels use [N, H, V, K] state layout. cuLA's + # Qwen3.5 wrapper uses [N, H, K, V], so pass the transposed view here. + initial_state_vk = initial_state.transpose(-1, -2).contiguous() + kwargs = dict( + q=q, + k=k, + v=v, + g=log_gate, + beta=beta, + initial_state=initial_state_vk, + initial_state_indices=initial_state_indices, + output_final_state=True, + scale=CONFIG.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + head_first=False, + ) + try: + sig = inspect.signature(chunk_gdr) + kwargs = {key: value for key, value in kwargs.items() if key in sig.parameters} + except (TypeError, ValueError): + pass + return chunk_gdr(**kwargs) + + +def _normalize_chunk_result(result): + if isinstance(result, tuple): + out = result[0] + state = result[-1] if len(result) >= 2 else None + return out, state + return result, None + + +def _state_to_cula_layout(state: torch.Tensor | None) -> torch.Tensor | None: + if state is None: + return None + return state.transpose(-1, -2).contiguous() + + +def print_header(device: torch.device, args: argparse.Namespace, baseline_sources: dict[str, str]) -> None: + print("Qwen3.5 prefill benchmark") + print(f" device: {torch.cuda.get_device_name(device)}") + print(f" dtype: {CONFIG.qkv_dtype}") + print(f" batch: {args.batch}") + print(f" seq lens: {args.seq_lens}") + print(f" warmup/rep: {args.warmup}/{args.rep}") + print(f" baselines: {baseline_sources or 'disabled/unavailable'}") + print() + print( + f"{'baseline':>8} {'B':>3} {'T':>7} {'layout_ms':>11} {'cula_kda_ms':>12} {'cula_total':>11} " + f"{'base_ms':>11} {'speedup':>9} {'rel_rms':>10} {'rel_max':>10}" + ) + print("-" * 108) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-lens", type=int, nargs="+", default=[128, 256, 512, 1024]) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rep", type=int, default=30) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--baseline", choices=["none", "fla", "sgl", "all"], default="sgl") + parser.add_argument("--sglang-path", type=pathlib.Path, default=None) + parser.add_argument("--skip-accuracy", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + device = torch.device("cuda") + + baselines: dict[str, Callable] = {} + baseline_sources: dict[str, str] = {} + if args.baseline in ("fla", "all"): + fla_chunk_gdr, fla_source_or_error = resolve_fla_chunk_gdr() + if fla_chunk_gdr is None: + print(f"Skipping FLA baseline: {fla_source_or_error}") + else: + baselines["fla"] = fla_chunk_gdr + baseline_sources["fla"] = fla_source_or_error + if args.baseline in ("sgl", "all"): + sgl_chunk_gdr, sgl_source_or_error = resolve_sgl_chunk_gdr(args.sglang_path) + if sgl_chunk_gdr is None: + print(f"Skipping SGLang baseline: {sgl_source_or_error}") + else: + baselines["sgl"] = sgl_chunk_gdr + baseline_sources["sgl"] = sgl_source_or_error + + print_header(device, args, baseline_sources) + + for seq_len in args.seq_lens: + q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, mixed_qkv_conv, a_flat, b_flat = make_inputs( + args.batch, + seq_len, + device=device, + seed=args.seed, + ) + initial_state_indices = torch.arange(args.batch, device=device, dtype=torch.int32) + + def layout_fn(): + return qwen35_layout_prefill(mixed_qkv_conv, a_flat, b_flat, backend="cudac") + + def cula_kda_fn(): + return run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state) + + layout_ms = benchmark_cuda_fn(layout_fn, warmup=args.warmup, rep=args.rep) + cula_kda_ms = benchmark_cuda_fn(cula_kda_fn, warmup=args.warmup, rep=args.rep) + cula_total_ms = layout_ms + cula_kda_ms + + if not baselines: + print( + f"{'none':>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_kda_ms:12.4f} {cula_total_ms:11.4f} " + f"{float('nan'):11.4f} {float('nan'):9.3f} {float('nan'):10.3e} {float('nan'):10.3e}" + ) + + for baseline_name, chunk_gdr in baselines.items(): + def baseline_fn(): + return run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices) + + rel_rms = float("nan") + rel_max = float("nan") + if not args.skip_accuracy: + out_cula, state_cula = cula_kda_fn() + out_base, state_base = _normalize_chunk_result(baseline_fn()) + state_base = _state_to_cula_layout(state_base) + torch.cuda.synchronize() + rel_rms, rel_max, _ = error_stats(out_base, out_cula) + if state_base is not None and tuple(state_base.shape) == tuple(state_cula.shape): + rel_rms_s, rel_max_s, _ = error_stats(state_base, state_cula) + rel_rms = max(rel_rms, rel_rms_s) + rel_max = max(rel_max, rel_max_s) + + base_ms = benchmark_cuda_fn(baseline_fn, warmup=args.warmup, rep=args.rep) + speedup = base_ms / cula_kda_ms if cula_kda_ms > 0 else float("inf") + print( + f"{baseline_name:>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_kda_ms:12.4f} {cula_total_ms:11.4f} " + f"{base_ms:11.4f} {speedup:9.3f} {rel_rms:10.3e} {rel_max:10.3e}" + ) + + del q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, initial_state_indices, mixed_qkv_conv, a_flat, b_flat + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index 59bc3e71..2e93cf0a 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -18,6 +18,7 @@ #include #include "qwen35/decode/qwen35_decode_common.cuh" +#include "qwen35/prefill/qwen35_prefill_common.cuh" #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) void @@ -158,6 +159,58 @@ qwen35_layout_scalar_kda_decode( cula::qwen35::decode::run_qwen35_layout_scalar_kda_decode(params); } +void +qwen35_scalar_kda_prefill( + at::Tensor q, + at::Tensor k, + at::Tensor v, + at::Tensor a, + at::Tensor b, + at::Tensor A_log, + at::Tensor dt_bias, + at::Tensor initial_state, + at::Tensor cu_seqlens, + at::Tensor out, + at::Tensor final_state) { + cula::qwen35::prefill::ScalarKdaPrefillParams params{ + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + }; + cula::qwen35::prefill::run_qwen35_scalar_kda_prefill(params); +} + +void +qwen35_layout_prefill( + at::Tensor mixed_qkv_conv, + at::Tensor a, + at::Tensor b, + at::Tensor q_rep, + at::Tensor k_rep, + at::Tensor v, + at::Tensor a_kernel, + at::Tensor b_kernel) { + cula::qwen35::prefill::LayoutPrefillParams params{ + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + }; + cula::qwen35::prefill::run_qwen35_layout_prefill(params); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "cuLA"; #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) @@ -171,4 +224,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("qwen35_layout_decode", &qwen35_layout_decode); m.def("qwen35_scalar_kda_decode", &qwen35_scalar_kda_decode); m.def("qwen35_layout_scalar_kda_decode", &qwen35_layout_scalar_kda_decode); + m.def("qwen35_layout_prefill", &qwen35_layout_prefill); + m.def("qwen35_scalar_kda_prefill", &qwen35_scalar_kda_prefill); } diff --git a/csrc/qwen35/prefill/qwen35_layout_prefill.cu b/csrc/qwen35/prefill/qwen35_layout_prefill.cu new file mode 100644 index 00000000..dbda9577 --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_layout_prefill.cu @@ -0,0 +1,127 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qwen35_layout_prefill_kernel.hpp" +#include "qwen35_prefill_common.cuh" + +#include +#include +#include +#include + +namespace cula::qwen35::prefill { + +namespace { + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); +} + +void check_rank_2(const at::Tensor& tensor, const char* name) { + TORCH_CHECK(tensor.dim() == 2, name, " must be rank 2, got rank ", tensor.dim(), "."); +} + +} // namespace + +void run_qwen35_layout_prefill(LayoutPrefillParams& params) { + const at::Tensor& mixed_qkv_conv = params.mixed_qkv_conv; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& q_rep = params.q_rep; + const at::Tensor& k_rep = params.k_rep; + const at::Tensor& v = params.v; + const at::Tensor& a_kernel = params.a_kernel; + const at::Tensor& b_kernel = params.b_kernel; + + TORCH_CHECK(mixed_qkv_conv.is_cuda(), "mixed_qkv_conv must be a CUDA tensor."); + const at::Device device = mixed_qkv_conv.device(); + + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(q_rep, "q_rep", device); + check_tensor_device(k_rep, "k_rep", device); + check_tensor_device(v, "v", device); + check_tensor_device(a_kernel, "a_kernel", device); + check_tensor_device(b_kernel, "b_kernel", device); + + TORCH_CHECK(mixed_qkv_conv.is_contiguous(), "mixed_qkv_conv must be contiguous."); + TORCH_CHECK(a.is_contiguous(), "a must be contiguous."); + TORCH_CHECK(b.is_contiguous(), "b must be contiguous."); + TORCH_CHECK(q_rep.is_contiguous(), "q_rep must be contiguous."); + TORCH_CHECK(k_rep.is_contiguous(), "k_rep must be contiguous."); + TORCH_CHECK(v.is_contiguous(), "v must be contiguous."); + TORCH_CHECK(a_kernel.is_contiguous(), "a_kernel must be contiguous."); + TORCH_CHECK(b_kernel.is_contiguous(), "b_kernel must be contiguous."); + + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == a.scalar_type() && + mixed_qkv_conv.scalar_type() == b.scalar_type() && + mixed_qkv_conv.scalar_type() == q_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == k_rep.scalar_type() && + mixed_qkv_conv.scalar_type() == v.scalar_type() && + mixed_qkv_conv.scalar_type() == a_kernel.scalar_type() && + mixed_qkv_conv.scalar_type() == b_kernel.scalar_type(), + "All layout prefill tensors must share the same dtype."); + TORCH_CHECK( + mixed_qkv_conv.scalar_type() == at::kHalf || mixed_qkv_conv.scalar_type() == at::kBFloat16, + "mixed_qkv_conv must be float16 or bfloat16."); + + check_rank_2(mixed_qkv_conv, "mixed_qkv_conv"); + check_rank_2(a, "a"); + check_rank_2(b, "b"); + + const int64_t token_count = mixed_qkv_conv.size(0); + TORCH_CHECK(mixed_qkv_conv.size(1) == kMixedQKVDim, "mixed_qkv_conv must be [N, 10240]."); + TORCH_CHECK(a.sizes() == at::IntArrayRef({token_count, kNumVHeads}), "a must be [N, 48]."); + TORCH_CHECK(b.sizes() == at::IntArrayRef({token_count, kNumVHeads}), "b must be [N, 48]."); + TORCH_CHECK( + q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimQK}), + "q_rep must be [N, 48, 128]."); + TORCH_CHECK(k_rep.sizes() == q_rep.sizes(), "k_rep must match q_rep shape."); + TORCH_CHECK(v.sizes() == q_rep.sizes(), "v must match q_rep shape."); + TORCH_CHECK(a_kernel.sizes() == a.sizes(), "a_kernel must match a shape."); + TORCH_CHECK(b_kernel.sizes() == b.sizes(), "b_kernel must match b shape."); + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + if (mixed_qkv_conv.scalar_type() == at::kHalf) { + kernel::launch_qwen35_layout_prefill_kernel( + stream, + mixed_qkv_conv.data_ptr(), + a.data_ptr(), + b.data_ptr(), + q_rep.data_ptr(), + k_rep.data_ptr(), + v.data_ptr(), + a_kernel.data_ptr(), + b_kernel.data_ptr(), + token_count); + } else { + kernel::launch_qwen35_layout_prefill_kernel( + stream, + mixed_qkv_conv.data_ptr(), + a.data_ptr(), + b.data_ptr(), + q_rep.data_ptr(), + k_rep.data_ptr(), + v.data_ptr(), + a_kernel.data_ptr(), + b_kernel.data_ptr(), + token_count); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::prefill diff --git a/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp new file mode 100644 index 00000000..b2e0adaf --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp @@ -0,0 +1,141 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "qwen35_prefill_common.cuh" + +#include +#include +#include +#include + +namespace cula::qwen35::prefill::kernel { + +using namespace cute; + +template +CUTE_DEVICE void copy_prefill_vec_contiguous( + scalar_t* __restrict__ dst, + const scalar_t* __restrict__ src) { + constexpr int kBytes = sizeof(scalar_t) * kVec; + if constexpr (kBytes == 16 || kBytes == 8) { + using VecType = cutlass::AlignedArray; + const auto dst_addr = reinterpret_cast(dst); + const auto src_addr = reinterpret_cast(src); + if ((dst_addr % alignof(VecType) == 0) && (src_addr % alignof(VecType) == 0)) { + *reinterpret_cast(dst) = *reinterpret_cast(src); + return; + } + } + +#pragma unroll + for (int i = 0; i < kVec; ++i) { + dst[i] = src[i]; + } +} + +template +__global__ void qwen35_layout_prefill_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + scalar_t* __restrict__ q_rep, + scalar_t* __restrict__ k_rep, + scalar_t* __restrict__ v_out, + scalar_t* __restrict__ a_kernel, + scalar_t* __restrict__ b_kernel, + int64_t token_count) { + static_assert(kNumVHeads % kNumQKHeads == 0); + static_assert(kHeadDimQK == kHeadDimV); + constexpr int kRepeatFactor = kNumVHeads / kNumQKHeads; + constexpr int kVec = 4; + static_assert(kHeadDimQK % kVec == 0); + + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + if (token_idx >= token_count || hv >= kNumVHeads) { + return; + } + + const int mapped_h = hv / kRepeatFactor; + + auto qk_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto hv_layout = make_layout( + make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})); + auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + + const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kMixedQKVDim; + const scalar_t* q_src_ptr = token_ptr; + const scalar_t* k_src_ptr = token_ptr + kQDim; + const scalar_t* v_src_ptr = token_ptr + kQDim + kKDim; + + scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; + scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; + scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kNumVHeads * kHeadDimV; + + for (int vec_idx = tid; vec_idx < kHeadDimQK / kVec; vec_idx += blockDim.x) { + const int d = vec_idx * kVec; + const int q_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int k_src_idx = crd2idx(make_coord(mapped_h, d), qk_src_layout); + const int v_src_idx = crd2idx(make_coord(hv, d), v_src_layout); + const int dst_idx = crd2idx(make_coord(hv, d), hv_layout); + + copy_prefill_vec_contiguous(q_dst_ptr + dst_idx, q_src_ptr + q_src_idx); + copy_prefill_vec_contiguous(k_dst_ptr + dst_idx, k_src_ptr + k_src_idx); + copy_prefill_vec_contiguous(v_dst_ptr + dst_idx, v_src_ptr + v_src_idx); + } + + if (tid == 0) { + const int head_idx = crd2idx(make_coord(hv), head_layout); + const int64_t token_head_offset = static_cast(token_idx) * kNumVHeads + head_idx; + a_kernel[token_head_offset] = a[token_head_offset]; + b_kernel[token_head_offset] = b[token_head_offset]; + } +} + +template +void launch_qwen35_layout_prefill_kernel( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t token_count) { + constexpr int kThreads = 32; + dim3 grid(kNumVHeads, static_cast(token_count), 1); + qwen35_layout_prefill_kernel<<>>( + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + token_count); +} + +} // namespace cula::qwen35::prefill::kernel diff --git a/csrc/qwen35/prefill/qwen35_prefill_common.cuh b/csrc/qwen35/prefill/qwen35_prefill_common.cuh new file mode 100644 index 00000000..0aac63b2 --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_prefill_common.cuh @@ -0,0 +1,60 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "qwen35/decode/qwen35_decode_common.cuh" + +#include + +namespace cula::qwen35::prefill { + +using decode::kHeadDimQK; +using decode::kHeadDimV; +using decode::kKDim; +using decode::kMixedQKVDim; +using decode::kNumQKHeads; +using decode::kNumVHeads; +using decode::kQDim; +using decode::kVDim; + +struct LayoutPrefillParams { + at::Tensor mixed_qkv_conv; // [N, 10240] + at::Tensor a; // [N, 48] + at::Tensor b; // [N, 48] + at::Tensor q_rep; // [N, 48, 128] + at::Tensor k_rep; // [N, 48, 128] + at::Tensor v; // [N, 48, 128] + at::Tensor a_kernel; // [N, 48] + at::Tensor b_kernel; // [N, 48] +}; + +struct ScalarKdaPrefillParams { + at::Tensor q; // [B, T, 48, 128] + at::Tensor k; // [B, T, 48, 128] + at::Tensor v; // [B, T, 48, 128] + at::Tensor a; // [B, T, 48] + at::Tensor b; // [B, T, 48] + at::Tensor A_log; // [48], float32 + at::Tensor dt_bias; // [48], float32 + at::Tensor initial_state; // [N, 48, 128, 128], float32, may be empty + at::Tensor cu_seqlens; // [N + 1], int32, may be empty + at::Tensor out; // [B, T, 48, 128] + at::Tensor final_state; // [N, 48, 128, 128], float32 +}; + +void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params); +void run_qwen35_layout_prefill(LayoutPrefillParams& params); + +} // namespace cula::qwen35::prefill diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu new file mode 100644 index 00000000..454d61ae --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu @@ -0,0 +1,171 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qwen35_prefill_common.cuh" +#include "qwen35_scalar_kda_prefill_kernel.hpp" + +#include +#include +#include +#include + +namespace cula::qwen35::prefill { + +namespace { + +void check_tensor_device(const at::Tensor& tensor, const char* name, const at::Device& device) { + if (tensor.defined() && tensor.numel() > 0) { + TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); + } +} + +void check_contiguous(const at::Tensor& tensor, const char* name) { + if (tensor.defined() && tensor.numel() > 0) { + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous."); + } +} + +} // namespace + +void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { + const at::Tensor& q = params.q; + const at::Tensor& k = params.k; + const at::Tensor& v = params.v; + const at::Tensor& a = params.a; + const at::Tensor& b = params.b; + const at::Tensor& A_log = params.A_log; + const at::Tensor& dt_bias = params.dt_bias; + const at::Tensor& initial_state = params.initial_state; + const at::Tensor& cu_seqlens = params.cu_seqlens; + const at::Tensor& out = params.out; + const at::Tensor& final_state = params.final_state; + + TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor."); + const at::Device device = q.device(); + + check_tensor_device(k, "k", device); + check_tensor_device(v, "v", device); + check_tensor_device(a, "a", device); + check_tensor_device(b, "b", device); + check_tensor_device(A_log, "A_log", device); + check_tensor_device(dt_bias, "dt_bias", device); + check_tensor_device(initial_state, "initial_state", device); + check_tensor_device(cu_seqlens, "cu_seqlens", device); + check_tensor_device(out, "out", device); + check_tensor_device(final_state, "final_state", device); + + check_contiguous(q, "q"); + check_contiguous(k, "k"); + check_contiguous(v, "v"); + check_contiguous(a, "a"); + check_contiguous(b, "b"); + check_contiguous(A_log, "A_log"); + check_contiguous(dt_bias, "dt_bias"); + check_contiguous(initial_state, "initial_state"); + check_contiguous(cu_seqlens, "cu_seqlens"); + check_contiguous(out, "out"); + check_contiguous(final_state, "final_state"); + + TORCH_CHECK( + q.scalar_type() == k.scalar_type() && q.scalar_type() == v.scalar_type() && + q.scalar_type() == a.scalar_type() && q.scalar_type() == b.scalar_type() && + q.scalar_type() == out.scalar_type(), + "q/k/v/a/b/out must share the same dtype."); + TORCH_CHECK(q.scalar_type() == at::kHalf || q.scalar_type() == at::kBFloat16, "q must be float16 or bfloat16."); + TORCH_CHECK(A_log.scalar_type() == at::kFloat, "A_log must be float32."); + TORCH_CHECK(dt_bias.scalar_type() == at::kFloat, "dt_bias must be float32."); + TORCH_CHECK(final_state.scalar_type() == at::kFloat, "final_state must be float32."); + TORCH_CHECK( + !initial_state.defined() || initial_state.numel() == 0 || initial_state.scalar_type() == at::kFloat, + "initial_state must be float32 when provided."); + TORCH_CHECK( + !cu_seqlens.defined() || cu_seqlens.numel() == 0 || cu_seqlens.scalar_type() == at::kInt, + "cu_seqlens must be int32 when provided."); + + TORCH_CHECK(q.dim() == 4, "q must be [B, T, 48, 128]."); + const int64_t B = q.size(0); + const int64_t T = q.size(1); + TORCH_CHECK( + q.sizes() == at::IntArrayRef({B, T, kNumVHeads, kHeadDimQK}), + "q must have shape [B, T, 48, 128]."); + TORCH_CHECK(k.sizes() == q.sizes(), "k must match q shape."); + TORCH_CHECK(v.sizes() == q.sizes(), "v must match q shape."); + TORCH_CHECK(a.dim() == 3 && a.sizes() == at::IntArrayRef({B, T, kNumVHeads}), "a must be [B, T, 48]."); + TORCH_CHECK(b.sizes() == a.sizes(), "b must match a shape."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == kNumVHeads, "A_log must be [48]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == kNumVHeads, "dt_bias must be [48]."); + TORCH_CHECK(out.sizes() == q.sizes(), "out must match q shape."); + + const bool is_varlen = cu_seqlens.defined() && cu_seqlens.numel() > 0; + const int64_t sequence_count = is_varlen ? cu_seqlens.numel() - 1 : B; + TORCH_CHECK(sequence_count > 0, "sequence_count must be positive."); + if (is_varlen) { + TORCH_CHECK(B == 1, "cu_seqlens mode expects flattened q/k/v with batch size 1."); + } + + TORCH_CHECK( + final_state.dim() == 4 && + final_state.sizes() == at::IntArrayRef({sequence_count, kNumVHeads, kHeadDimQK, kHeadDimV}), + "final_state must be [N, 48, 128, 128]."); + const bool has_initial_state = initial_state.defined() && initial_state.numel() > 0; + if (has_initial_state) { + TORCH_CHECK(initial_state.sizes() == final_state.sizes(), "initial_state must match final_state shape."); + } + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + + if (q.scalar_type() == at::kHalf) { + kernel::launch_qwen35_scalar_kda_prefill_kernel( + stream, + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + a.data_ptr(), + b.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + is_varlen ? cu_seqlens.data_ptr() : nullptr, + out.data_ptr(), + final_state.data_ptr(), + static_cast(B), + static_cast(T), + static_cast(sequence_count), + is_varlen, + has_initial_state); + } else { + kernel::launch_qwen35_scalar_kda_prefill_kernel( + stream, + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + a.data_ptr(), + b.data_ptr(), + A_log.data_ptr(), + dt_bias.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + is_varlen ? cu_seqlens.data_ptr() : nullptr, + out.data_ptr(), + final_state.data_ptr(), + static_cast(B), + static_cast(T), + static_cast(sequence_count), + is_varlen, + has_initial_state); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::prefill diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp new file mode 100644 index 00000000..f0aba77f --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp @@ -0,0 +1,273 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "qwen35_prefill_common.cuh" + +#include +#include + +namespace cula::qwen35::prefill::kernel { + +using namespace cute; + +template +struct Qwen35ScalarKdaPrefillKernel { + static constexpr int kThreads = 128; + static constexpr int kHeadDim = kHeadDimQK; + static constexpr int kVTile = 8; + static constexpr int kNumVTiles = kHeadDimV / kVTile; + + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kHeadDimV % kVTile == 0); + + struct SharedStorage { + float scratch[kThreads]; + }; + + static dim3 block_shape() { + return dim3(kThreads, 1, 1); + } + + CUTE_HOST_DEVICE static auto make_v_work_tiles(int sequence_count) { + auto problem_layout = make_layout( + make_shape(Int{}, Int{}, sequence_count), + make_stride(Int<1>{}, Int{}, Int{})); + return zipped_divide(problem_layout, make_shape(Int{}, Int<1>{}, Int<1>{})); + } + + static dim3 grid_shape(int sequence_count) { + auto v_work_tiles = make_v_work_tiles(sequence_count); + return dim3(static_cast(size<1>(v_work_tiles)), 1, 1); + } + + CUTE_DEVICE static float load_as_float(scalar_t value) { + return static_cast(value); + } + + CUTE_DEVICE static scalar_t cast_output(float value) { + return static_cast(value); + } + + CUTE_DEVICE static float softplus(float x) { + return x > 20.0f ? x : log1pf(expf(x)); + } + + CUTE_DEVICE static float block_sum(float value, SharedStorage& storage, int tid) { + storage.scratch[tid] = value; + __syncthreads(); + + for (int stride = kThreads / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + storage.scratch[tid] += storage.scratch[tid + stride]; + } + __syncthreads(); + } + return storage.scratch[0]; + } + + CUTE_DEVICE static void run_device( + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ initial_state, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int sequence_count, + bool is_varlen, + bool has_initial_state, + SharedStorage& storage) { + auto v_work_tiles = make_v_work_tiles(sequence_count); + auto work_layout = make_layout(get<1>(v_work_tiles.shape()), LayoutLeft{}); + auto work_coord = work_layout.get_hier_coord(static_cast(blockIdx.x)); + const int v_tile_idx = static_cast(get<0>(work_coord)); + const int hv = static_cast(get<1>(work_coord)); + const int seq_idx = static_cast(get<2>(work_coord)); + const int v_base = v_tile_idx * kVTile; + const int tid = static_cast(threadIdx.x); + + if (hv >= kNumVHeads || seq_idx >= sequence_count) { + return; + } + + const int token_begin = is_varlen ? static_cast(cu_seqlens[seq_idx]) : seq_idx * seq_len; + const int token_end = is_varlen ? static_cast(cu_seqlens[seq_idx + 1]) : token_begin + seq_len; + const int state_base = ((seq_idx * kNumVHeads + hv) * kHeadDimQK) * kHeadDimV; + + const int kk = tid; + float state_vals[kVTile]; + +#pragma unroll + for (int lane = 0; lane < kVTile; ++lane) { + const int v_row = v_base + lane; + const int state_off = state_base + kk * kHeadDimV + v_row; + state_vals[lane] = 0.0f; + if (kk < kHeadDimQK && v_row < kHeadDimV) { + state_vals[lane] = has_initial_state ? initial_state[state_off] : 0.0f; + } + } + __syncthreads(); + + const float scale = rsqrtf(static_cast(kHeadDimQK)); + const float exp_A = expf(A_log[hv]); + const float dt = dt_bias[hv]; + + for (int token = token_begin; token < token_end; ++token) { + const int local_t = is_varlen ? token : token - token_begin; + const int qkv_base = ((token * kNumVHeads + hv) * kHeadDimQK); + const int gate_base = token * kNumVHeads + hv; + + const float q_val = kk < kHeadDimQK ? load_as_float(q[qkv_base + kk]) : 0.0f; + const float k_val = kk < kHeadDimQK ? load_as_float(k[qkv_base + kk]) : 0.0f; + const float q_norm_sq = block_sum(q_val * q_val, storage, tid); + const float k_norm_sq = block_sum(k_val * k_val, storage, tid); + const float q_rnorm = rsqrtf(fmaxf(q_norm_sq, 1.0e-20f)) * scale; + const float k_rnorm = rsqrtf(fmaxf(k_norm_sq, 1.0e-20f)); + + const float decay = expf(-exp_A * softplus(load_as_float(a[gate_base]) + dt)); + const float beta = 1.0f / (1.0f + expf(-load_as_float(b[gate_base]))); + + const float k_norm = k_val * k_rnorm; + const float q_norm = q_val * q_rnorm; + +#pragma unroll + for (int lane = 0; lane < kVTile; ++lane) { + const int v_row = v_base + lane; + if (v_row < kHeadDimV) { + const float proj_partial = kk < kHeadDimQK ? state_vals[lane] * k_norm : 0.0f; + const float proj = block_sum(proj_partial, storage, tid); + + const float v_val = load_as_float(v[qkv_base + v_row]); + const float v_new = beta * (v_val - decay * proj); + + float out_partial = 0.0f; + if (kk < kHeadDimQK) { + const float state_new = decay * state_vals[lane] + k_norm * v_new; + state_vals[lane] = state_new; + out_partial = state_new * q_norm; + } + const float out_acc = block_sum(out_partial, storage, tid); + + if (tid == 0) { + const int out_off = + (((is_varlen ? 0 : seq_idx) * seq_len + local_t) * kNumVHeads + hv) * kHeadDimV + v_row; + out[out_off] = cast_output(out_acc); + } + } + } + __syncthreads(); + } + +#pragma unroll + for (int lane = 0; lane < kVTile; ++lane) { + const int v_row = v_base + lane; + if (kk < kHeadDimQK && v_row < kHeadDimV) { + const int state_off = state_base + kk * kHeadDimV + v_row; + final_state[state_off] = state_vals[lane]; + } + } + + (void)batch_size; + } +}; + +template +__global__ void qwen35_scalar_kda_prefill_kernel( + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ initial_state, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + __shared__ typename Qwen35ScalarKdaPrefillKernel::SharedStorage storage; + Qwen35ScalarKdaPrefillKernel::run_device( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + batch_size, + seq_len, + sequence_count, + is_varlen, + has_initial_state, + storage); +} + +template +void launch_qwen35_scalar_kda_prefill_kernel( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + const auto grid = Qwen35ScalarKdaPrefillKernel::grid_shape(sequence_count); + const auto block = Qwen35ScalarKdaPrefillKernel::block_shape(); + qwen35_scalar_kda_prefill_kernel<<>>( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + batch_size, + seq_len, + sequence_count, + is_varlen, + has_initial_state); +} + +} // namespace cula::qwen35::prefill::kernel diff --git a/cula/ops/qwen35_conv1d_prefill.py b/cula/ops/qwen35_conv1d_prefill.py index e153460b..3a32b1e0 100644 --- a/cula/ops/qwen35_conv1d_prefill.py +++ b/cula/ops/qwen35_conv1d_prefill.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CuTe DSL placeholder for Qwen3.5 depthwise causal conv1d prefill.""" +"""Qwen3.5 depthwise causal conv1d prefill wrapper.""" from __future__ import annotations @@ -24,13 +24,76 @@ def qwen35_conv1d_prefill( weight: torch.Tensor, *, activation: str = "silu", -) -> torch.Tensor: + cu_seqlens: torch.Tensor | None = None, + output_final_state: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Depthwise causal conv1d over a full sequence. Expected shapes: - - x: [B, C, S] + - x: [B, C, S] or flattened [T, C] - weight: [C, 1, 4] or [C, 4] """ - del x, weight, activation - raise NotImplementedError("Qwen3.5 conv1d prefill kernel is not implemented yet.") + if activation != "silu": + raise ValueError(f"Unsupported activation={activation}") + if weight.ndim == 3: + if weight.shape[1] != 1: + raise ValueError(f"weight must be [C,1,K] or [C,K], got {tuple(weight.shape)}") + weight_2d = weight[:, 0, :] + elif weight.ndim == 2: + weight_2d = weight + else: + raise ValueError(f"weight must be [C,1,K] or [C,K], got {tuple(weight.shape)}") + + kernel_size = weight_2d.shape[1] + + def _conv_one(seq: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # seq: [S, C] + if seq.ndim != 2 or seq.shape[1] != weight_2d.shape[0]: + raise ValueError(f"sequence must be [S,C={weight_2d.shape[0]}], got {tuple(seq.shape)}") + seq_f = seq.float() + weight_f = weight_2d.float() + out = torch.empty_like(seq) + for t in range(seq.shape[0]): + acc = torch.zeros(seq.shape[1], device=seq.device, dtype=torch.float32) + for kk in range(kernel_size): + src_t = t - (kernel_size - 1 - kk) + if src_t >= 0: + acc = acc + seq_f[src_t] * weight_f[:, kk] + out[t] = torch.nn.functional.silu(acc).to(seq.dtype) + + state = torch.zeros(seq.shape[1], kernel_size, device=seq.device, dtype=seq.dtype) + take = min(kernel_size, seq.shape[0]) + if take > 0: + state[:, kernel_size - take :] = seq[-take:].transpose(0, 1) + return out, state + + if x.ndim == 3: + # Public op shape follows the Qwen conv convention [B, C, S]. + if x.shape[1] != weight_2d.shape[0]: + raise ValueError(f"x channel dim must match weight, got x={tuple(x.shape)} weight={tuple(weight_2d.shape)}") + y = torch.empty_like(x) + states = torch.empty(x.shape[0], x.shape[1], kernel_size, device=x.device, dtype=x.dtype) + for bidx in range(x.shape[0]): + y_b, state_b = _conv_one(x[bidx].transpose(0, 1).contiguous()) + y[bidx] = y_b.transpose(0, 1).contiguous() + states[bidx] = state_b + return (y, states) if output_final_state else y + + if x.ndim == 2: + if cu_seqlens is None: + y, state = _conv_one(x) + return (y, state.unsqueeze(0)) if output_final_state else y + if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + y = torch.empty_like(x) + states = torch.empty(cu_seqlens.numel() - 1, x.shape[1], kernel_size, device=x.device, dtype=x.dtype) + for sidx in range(cu_seqlens.numel() - 1): + start = int(cu_seqlens[sidx].item()) + end = int(cu_seqlens[sidx + 1].item()) + y_s, state_s = _conv_one(x[start:end]) + y[start:end] = y_s + states[sidx] = state_s + return (y, states) if output_final_state else y + + raise ValueError(f"x must be [B,C,S] or [T,C], got {tuple(x.shape)}") diff --git a/cula/ops/qwen35_layout_prefill.py b/cula/ops/qwen35_layout_prefill.py new file mode 100644 index 00000000..64257d22 --- /dev/null +++ b/cula/ops/qwen35_layout_prefill.py @@ -0,0 +1,97 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3.5 layout prefill wrapper.""" + +from __future__ import annotations + +import torch + +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig, infer_local_config + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + + +def qwen35_layout_prefill_reference( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + local_key_dim, _, local_num_k_heads = infer_local_config( + mixed_qkv_conv.shape[1], + local_num_v_heads, + config=config, + ) + + q_end = local_key_dim + k_end = q_end + local_key_dim + q = mixed_qkv_conv[:, :q_end].view(tokens, local_num_k_heads, config.head_k_dim) + k = mixed_qkv_conv[:, q_end:k_end].view(tokens, local_num_k_heads, config.head_k_dim) + v = mixed_qkv_conv[:, k_end:].view(tokens, local_num_v_heads, config.head_v_dim) + + repeat_factor = local_num_v_heads // local_num_k_heads + q_rep = q.repeat_interleave(repeat_factor, dim=1).contiguous() + k_rep = k.repeat_interleave(repeat_factor, dim=1).contiguous() + return q_rep, k_rep, v.contiguous(), a.contiguous(), b.contiguous() + + +def qwen35_layout_prefill( + mixed_qkv_conv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_layout_prefill") + and mixed_qkv_conv.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_layout_prefill is not available.") + + if use_cudac: + tokens = mixed_qkv_conv.shape[0] + local_num_v_heads = a.shape[1] + if local_num_v_heads != 48: + raise ValueError(f"backend='cudac' currently expects Qwen3.5 HV=48, got {local_num_v_heads}") + q_rep = torch.empty(tokens, local_num_v_heads, config.head_k_dim, device=mixed_qkv_conv.device, dtype=mixed_qkv_conv.dtype) + k_rep = torch.empty_like(q_rep) + v = torch.empty(tokens, local_num_v_heads, config.head_v_dim, device=mixed_qkv_conv.device, dtype=mixed_qkv_conv.dtype) + a_kernel = torch.empty_like(a) + b_kernel = torch.empty_like(b) + cula_cuda.qwen35_layout_prefill( + mixed_qkv_conv.contiguous(), + a.contiguous(), + b.contiguous(), + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + ) + return q_rep, k_rep, v, a_kernel, b_kernel + + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + return qwen35_layout_prefill_reference(mixed_qkv_conv, a, b, config=config) diff --git a/cula/ops/qwen35_scalar_kda_prefill.py b/cula/ops/qwen35_scalar_kda_prefill.py index 2fda06cb..ae494356 100644 --- a/cula/ops/qwen35_scalar_kda_prefill.py +++ b/cula/ops/qwen35_scalar_kda_prefill.py @@ -12,12 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CuTe DSL placeholder for Qwen3.5 scalar-gated KDA prefill.""" +"""Qwen3.5 scalar-gated KDA prefill wrapper.""" from __future__ import annotations import torch +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + def qwen35_scalar_kda_prefill( q: torch.Tensor, @@ -30,8 +35,109 @@ def qwen35_scalar_kda_prefill( *, initial_state: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, + backend: str = "auto", ) -> tuple[torch.Tensor, torch.Tensor | None]: """Chunked scalar-gated delta-rule prefill for Qwen3.5.""" - del q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens - raise NotImplementedError("Qwen3.5 scalar-gated KDA prefill kernel is not implemented yet.") + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError(f"q/k/v must be 4D [B,T,HV,D], got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + if q.shape != k.shape or q.shape != v.shape: + raise ValueError(f"q/k/v must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + B, T, HV, K = q.shape + if K != 128 or v.shape[-1] != 128: + raise ValueError(f"Qwen3.5 prefill expects K=V=128, got q={tuple(q.shape)} v={tuple(v.shape)}") + if a.ndim == 2: + a = a.unsqueeze(0) + if b.ndim == 2: + b = b.unsqueeze(0) + if a.shape != (B, T, HV) or b.shape != (B, T, HV): + raise ValueError(f"a/b must be [B,T,HV], got a={tuple(a.shape)} b={tuple(b.shape)} expected={(B, T, HV)}") + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + if cu_seqlens is not None: + if B != 1: + raise ValueError("cu_seqlens mode expects flattened q/k/v with batch size 1") + if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + if initial_state is not None and initial_state.shape[1:] != (HV, K, K): + raise ValueError(f"initial_state must be [N,HV,128,128], got {tuple(initial_state.shape)}") + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_scalar_kda_prefill") + and q.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_prefill is not available.") + + if use_cudac: + if HV != 48: + raise ValueError(f"backend='cudac' currently expects Qwen3.5 HV=48, got {HV}") + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + out = torch.empty_like(v) + final_state = torch.empty(state_count, HV, K, K, device=q.device, dtype=torch.float32) + initial_state_arg = ( + torch.empty(0, device=q.device, dtype=torch.float32) + if initial_state is None + else initial_state.contiguous() + ) + cu_seqlens_arg = ( + torch.empty(0, device=q.device, dtype=torch.int32) + if cu_seqlens is None + else cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + ) + cula_cuda.qwen35_scalar_kda_prefill( + q.contiguous(), + k.contiguous(), + v.contiguous(), + a.contiguous(), + b.contiguous(), + A_log.contiguous(), + dt_bias.contiguous(), + initial_state_arg, + cu_seqlens_arg, + out, + final_state, + ) + return out, final_state + + if backend not in ("auto", "reference"): + raise ValueError(f"Unsupported backend={backend}") + + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + state = ( + torch.zeros(state_count, HV, K, K, device=q.device, dtype=torch.float32) + if initial_state is None + else initial_state.float().clone() + ) + out = torch.empty_like(v) + q_f = torch.nn.functional.normalize(q.float(), dim=-1) * (K**-0.5) + k_f = torch.nn.functional.normalize(k.float(), dim=-1) + v_f = v.float() + a_f = a.float() + b_f = b.float() + A_log_f = A_log.float() + dt_bias_f = dt_bias.float() + + def _run_sequence(batch_idx: int, state_idx: int, start: int, end: int) -> None: + for t in range(start, end): + for hv in range(HV): + state_kv = state[state_idx, hv] + decay = torch.exp(-torch.exp(A_log_f[hv]) * torch.nn.functional.softplus(a_f[batch_idx, t, hv] + dt_bias_f[hv])) + beta = torch.sigmoid(b_f[batch_idx, t, hv]) + k_vec = k_f[batch_idx, t, hv] + q_vec = q_f[batch_idx, t, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v_f[batch_idx, t, hv] - proj) + state_kv_new = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[batch_idx, t, hv] = (state_kv_new.transpose(0, 1) @ q_vec).to(out.dtype) + state[state_idx, hv] = state_kv_new + + if cu_seqlens is None: + for bidx in range(B): + _run_sequence(bidx, bidx, 0, T) + else: + for sidx in range(state_count): + _run_sequence(0, sidx, int(cu_seqlens[sidx].item()), int(cu_seqlens[sidx + 1].item())) + return out, state diff --git a/cula/qwen35/runtime.py b/cula/qwen35/runtime.py index e2b5323e..6be6371e 100644 --- a/cula/qwen35/runtime.py +++ b/cula/qwen35/runtime.py @@ -32,12 +32,15 @@ validate_state_tensors, ) from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_update +from cula.ops.qwen35_conv1d_prefill import qwen35_conv1d_prefill from cula.ops.qwen35_layout_decode import qwen35_layout_decode +from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill from cula.ops.qwen35_scalar_kda_decode import ( has_qwen35_layout_scalar_kda_decode_cudac, qwen35_layout_scalar_kda_decode, qwen35_scalar_kda_decode, ) +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill _stream_cache: dict[tuple[str, int], object] = {} @@ -125,7 +128,7 @@ def qwen35_linear_attention_decode_reference( activation="silu", backend="reference", ) - q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode( + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_prefill( conv_out, a, b, @@ -158,19 +161,75 @@ def qwen35_linear_attention_prefill( cu_seqlens: torch.Tensor | None = None, recurrent_state: torch.Tensor | None = None, conv_state: torch.Tensor | None = None, + backend: str = "auto", ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: """Qwen3.5 prefill wrapper. - This is a thin runtime boundary. The underlying CuTe kernels are added in - dedicated `cula.ops.qwen35_*` modules. + Args: + mixed_qkv: flattened [tokens, local_conv_dim] + a, b: [tokens, local_num_v_heads] + conv_weight: [local_conv_dim, 1, 4] or [local_conv_dim, 4] + A_log, dt_bias: [local_num_v_heads] + cu_seqlens: optional int32 sequence offsets for flattened input + recurrent_state: optional initial recurrent state [num_sequences, HV, 128, 128] + + Returns: + - core_attn_out_flat: [tokens, local_value_dim] + - final conv_state: [num_sequences, local_conv_dim, 4] + - final recurrent_state: [num_sequences, local_num_v_heads, 128, 128] """ - del conv_weight, A_log, dt_bias, cu_seqlens validate_mixed_qkv(mixed_qkv, config) validate_scalar_gate_inputs(a, b, config) validate_state_tensors(conv_state, recurrent_state, config) - _get_cached_stream(mixed_qkv.device) - raise NotImplementedError("Qwen3.5 prefill kernel path is not implemented yet.") + if mixed_qkv.is_cuda: + _get_cached_stream(mixed_qkv.device) + if conv_state is not None: + raise NotImplementedError("Qwen3.5 prefill with non-empty conv_state is not implemented yet.") + if mixed_qkv.shape[0] != a.shape[0]: + raise ValueError(f"Token dimension mismatch, got mixed_qkv={tuple(mixed_qkv.shape)} a={tuple(a.shape)}") + if A_log.ndim != 1 or dt_bias.ndim != 1 or A_log.shape != dt_bias.shape: + raise ValueError(f"A_log and dt_bias must be matching 1D tensors, got {tuple(A_log.shape)} and {tuple(dt_bias.shape)}") + if cu_seqlens is not None and (cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32): + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + + tokens = mixed_qkv.shape[0] + local_num_v_heads = a.shape[1] + _, local_value_dim, _ = infer_local_config( + mixed_qkv.shape[1], + local_num_v_heads, + config=config, + ) + if A_log.numel() != local_num_v_heads: + raise ValueError(f"A_log must match local_num_v_heads={local_num_v_heads}, got {A_log.numel()}") + + conv_out, conv_state_out = qwen35_conv1d_prefill( + mixed_qkv, + conv_weight, + activation="silu", + cu_seqlens=cu_seqlens, + output_final_state=True, + ) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_prefill( + conv_out, + a, + b, + config=config, + backend="reference" if backend == "reference" else "auto", + ) + core_attn_out, recurrent_state_out = qwen35_scalar_kda_prefill( + q=q_rep.unsqueeze(0).contiguous(), + k=k_rep.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a_kernel.unsqueeze(0).contiguous(), + b=b_kernel.unsqueeze(0).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + initial_state=recurrent_state, + cu_seqlens=cu_seqlens, + backend=backend, + ) + return core_attn_out.reshape(tokens, local_value_dim), conv_state_out, recurrent_state_out def qwen35_linear_attention_decode( diff --git a/setup.py b/setup.py index 411401ea..8ee1bc55 100644 --- a/setup.py +++ b/setup.py @@ -150,6 +150,8 @@ def get_nvcc_thread_args(): "csrc/qwen35/decode/qwen35_conv1d_decode.cu", "csrc/qwen35/decode/qwen35_layout_decode.cu", "csrc/qwen35/decode/qwen35_scalar_kda_decode.cu", + "csrc/qwen35/prefill/qwen35_layout_prefill.cu", + "csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu", ] if not DISABLE_SM100 or not DISABLE_SM103: cuda_sources.extend( diff --git a/tests/test_qwen35_prefill.py b/tests/test_qwen35_prefill.py new file mode 100644 index 00000000..a1e0b640 --- /dev/null +++ b/tests/test_qwen35_prefill.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pathlib +import sys + +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.ops.qwen35_conv1d_prefill import qwen35_conv1d_prefill +from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill, qwen35_layout_prefill_reference +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill +from cula.qwen35.common import Qwen35LinearAttentionConfig +from cula.qwen35.runtime import qwen35_linear_attention_prefill + + +def _manual_scalar_prefill(q, k, v, a, b, A_log, dt_bias, initial_state=None, cu_seqlens=None): + B, T, HV, K = q.shape + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + state = torch.zeros(state_count, HV, K, K, device=q.device, dtype=torch.float32) + if initial_state is not None: + state = initial_state.float().clone() + out = torch.empty_like(v) + q_f = torch.nn.functional.normalize(q.float(), dim=-1) * (K**-0.5) + k_f = torch.nn.functional.normalize(k.float(), dim=-1) + + def run_seq(batch_idx, state_idx, start, end): + for t in range(start, end): + for hv in range(HV): + state_kv = state[state_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv].float()) * torch.nn.functional.softplus(a[batch_idx, t, hv].float() + dt_bias[hv].float())) + beta = torch.sigmoid(b[batch_idx, t, hv].float()) + k_vec = k_f[batch_idx, t, hv] + q_vec = q_f[batch_idx, t, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v[batch_idx, t, hv].float() - proj) + state_new = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[batch_idx, t, hv] = (state_new.transpose(0, 1) @ q_vec).to(out.dtype) + state[state_idx, hv] = state_new + + if cu_seqlens is None: + for batch_idx in range(B): + run_seq(batch_idx, batch_idx, 0, T) + else: + for state_idx in range(state_count): + run_seq(0, state_idx, int(cu_seqlens[state_idx].item()), int(cu_seqlens[state_idx + 1].item())) + return out, state + + +def test_qwen35_scalar_kda_prefill_reference_matches_manual(): + torch.manual_seed(0) + B, T, HV, K = 2, 3, 2, 128 + q = torch.randn(B, T, HV, K, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(B, T, HV, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, dtype=torch.bfloat16) + A_log = -torch.rand(HV, dtype=torch.float32) + dt_bias = torch.randn(HV, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, dtype=torch.float32) * 0.01 + + out_ref, state_ref = _manual_scalar_prefill(q, k, v, a, b, A_log, dt_bias, initial_state) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + + torch.testing.assert_close(out.float(), out_ref.float(), atol=1e-3, rtol=1e-3) + torch.testing.assert_close(state, state_ref, atol=1e-4, rtol=1e-4) + + +def test_qwen35_scalar_kda_prefill_varlen_reference_matches_manual(): + torch.manual_seed(1) + T, HV, K = 4, 2, 128 + q = torch.randn(1, T, HV, K, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(1, T, HV, dtype=torch.bfloat16) + b = torch.randn(1, T, HV, dtype=torch.bfloat16) + A_log = -torch.rand(HV, dtype=torch.float32) + dt_bias = torch.randn(HV, dtype=torch.float32) * 0.1 + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + + out_ref, state_ref = _manual_scalar_prefill(q, k, v, a, b, A_log, dt_bias, cu_seqlens=cu_seqlens) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + cu_seqlens=cu_seqlens, + backend="reference", + ) + + torch.testing.assert_close(out.float(), out_ref.float(), atol=1e-3, rtol=1e-3) + torch.testing.assert_close(state, state_ref, atol=1e-4, rtol=1e-4) + + +def test_qwen35_conv1d_prefill_flattened_state(): + x = torch.arange(5 * 3, dtype=torch.bfloat16).reshape(5, 3) + weight = torch.ones(3, 4, dtype=torch.bfloat16) + cu_seqlens = torch.tensor([0, 2, 5], dtype=torch.int32) + + y, state = qwen35_conv1d_prefill(x, weight, cu_seqlens=cu_seqlens, output_final_state=True) + + assert y.shape == x.shape + assert state.shape == (2, 3, 4) + torch.testing.assert_close(state[0, :, -2:], x[:2].transpose(0, 1)) + torch.testing.assert_close(state[1, :, -3:], x[2:5].transpose(0, 1)) + + +def test_qwen35_layout_prefill_reference(): + torch.manual_seed(2) + config = Qwen35LinearAttentionConfig(num_k_heads=1, num_v_heads=2) + tokens = 3 + mixed_qkv = torch.randn(tokens, config.conv_dim, dtype=torch.bfloat16) + a = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + b = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + + ref = qwen35_layout_prefill_reference(mixed_qkv, a, b, config=config) + out = qwen35_layout_prefill(mixed_qkv, a, b, config=config, backend="reference") + + for out_tensor, ref_tensor in zip(out, ref, strict=True): + assert torch.equal(out_tensor, ref_tensor) + + +def test_qwen35_linear_attention_prefill_reference_shapes(): + torch.manual_seed(2) + config = Qwen35LinearAttentionConfig(num_k_heads=1, num_v_heads=2) + tokens = 3 + mixed_qkv = torch.randn(tokens, config.conv_dim, dtype=torch.bfloat16) + a = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + b = torch.randn(tokens, config.num_v_heads, dtype=torch.bfloat16) + conv_weight = torch.randn(config.conv_dim, config.conv_kernel_size, dtype=torch.bfloat16) + A_log = -torch.rand(config.num_v_heads, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, dtype=torch.float32) * 0.1 + cu_seqlens = torch.tensor([0, 2, 3], dtype=torch.int32) + + out, conv_state, recurrent_state = qwen35_linear_attention_prefill( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + cu_seqlens=cu_seqlens, + backend="reference", + ) + + assert out.shape == (tokens, config.value_dim) + assert conv_state.shape == (2, config.conv_dim, config.conv_kernel_size) + assert recurrent_state.shape == (2, config.num_v_heads, config.head_k_dim, config.head_v_dim) From 3f62c26df8e4d651ff0a8ba7cb24bf3dd8ca1216 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Thu, 11 Jun 2026 08:25:56 +0000 Subject: [PATCH 09/35] Add Qwen3.5 fused prefill benchmark path --- benchmarks/bench_qwen35_prefill.py | 177 +++++++++++++++-- csrc/api/pybind.cu | 6 + csrc/qwen35/prefill/qwen35_prefill_common.cuh | 6 + .../qwen35_scalar_kda_prefill_kernel.hpp | 6 +- .../prefill/sm90/qwen35_chunk_prefill_sm90.cu | 180 ++++++++++++++++++ .../sm90/qwen35_chunk_prefill_traits_sm90.hpp | 112 +++++++++++ cula/kda/__init__.py | 16 +- cula/kda/blackwell_fused_fwd.py | 61 +++++- cula/ops/kda_fully_fused_sm100_wip.py | 16 +- cula/ops/qwen35_fused_kda_prefill.py | 132 +++++++++++++ cula/utils.py | 4 +- setup.py | 1 + tests/test_qwen35_prefill.py | 121 ++++++++++++ 13 files changed, 801 insertions(+), 37 deletions(-) create mode 100644 csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu create mode 100644 csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_traits_sm90.hpp create mode 100644 cula/ops/qwen35_fused_kda_prefill.py diff --git a/benchmarks/bench_qwen35_prefill.py b/benchmarks/bench_qwen35_prefill.py index 5e7dbe7b..ef470b7f 100644 --- a/benchmarks/bench_qwen35_prefill.py +++ b/benchmarks/bench_qwen35_prefill.py @@ -17,13 +17,18 @@ Reports: - layout: cuLA Qwen3.5 prefill layout split/repeat kernel - - scalar_kda: cuLA Qwen3.5 scalar-gated KDA prefill kernel + - cula_qk: cuLA Qwen3.5 TMA/WGMMA-or-UMMA QK chunk debug kernel + - cula_fused: cuLA generic fused KDA core through a Qwen3.5 scalar-gate adapter - fla_gdr: optional FLA chunk_gated_delta_rule baseline - sgl_gdr: optional SGLang vendored Triton chunk_gated_delta_rule baseline Baselines are optional. SGLang Qwen3.5 prefill uses the same chunked gated delta rule family in its Triton GDN kernel; decode uses a recurrent packed kernel instead. + +Note: cula_qk currently benchmarks the TMA tensor-core Q @ K^T subpath only, +not the full gated-delta prefill recurrence. Its output is [B,48,T,T], so long +sequence lengths have quadratic memory cost. """ from __future__ import annotations @@ -37,13 +42,23 @@ from collections.abc import Callable import torch +import torch.nn.functional as F ROOT = pathlib.Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT)) from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill +from cula.ops.qwen35_fused_kda_prefill import qwen35_fused_kda_prefill from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as CONFIG +from cula.utils import get_kda_fused_fwd + +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + +RCP_LN2 = 1.4426950408889634 def benchmark_cuda_fn(fn: Callable[[], object], *, warmup: int, rep: int) -> float: @@ -129,6 +144,13 @@ def make_inputs(batch: int, seq_len: int, *, device: torch.device, seed: int): return q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, mixed_qkv_conv, a_flat, b_flat +def run_cula_chunk_qk(q, k, out): + if cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): + raise RuntimeError("cula.cudac.qwen35_chunk_qk_prefill_sm90 is not available. Rebuild the CUDA extension.") + cula_cuda.qwen35_chunk_qk_prefill_sm90(q.contiguous(), k.contiguous(), out) + return out + + def run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state): return qwen35_scalar_kda_prefill( q, @@ -143,6 +165,54 @@ def run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state): ) +def run_cula_fused(q, k, v, a, b, A_log, dt_bias, initial_state): + return qwen35_fused_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + ) + + +def prepare_cula_fused_core_inputs(q, k, a, b, A_log, dt_bias, initial_state): + B, T, HV, K = q.shape + q_norm = F.normalize(q.float(), dim=-1).to(q.dtype).contiguous() + k_norm = F.normalize(k.float(), dim=-1).to(k.dtype).contiguous() + log_gate_scalar = -torch.exp(A_log.float()).view(1, 1, HV, 1) * F.softplus( + a.float().unsqueeze(-1) + dt_bias.float().view(1, 1, HV, 1) + ) + log_gate = log_gate_scalar.expand(B, T, HV, K).contiguous() + chunks = [] + for chunk_start in range(0, T, 64): + chunks.append(log_gate[:, chunk_start : chunk_start + 64].cumsum(dim=1) * RCP_LN2) + log_gate_cumsum = torch.cat(chunks, dim=1).contiguous() + beta = torch.sigmoid(b.float()).contiguous() + initial_state_vk = initial_state.float().transpose(-1, -2).contiguous() + return q_norm, k_norm, log_gate_cumsum, beta, initial_state_vk + + +def run_cula_fused_core(q_norm, k_norm, v, log_gate_cumsum, beta, initial_state_vk): + fused_kda_prefill = get_kda_fused_fwd(q_norm.device) + return fused_kda_prefill( + q=q_norm, + k=k_norm, + v=v.contiguous(), + g=log_gate_cumsum, + beta=beta, + scale=CONFIG.head_k_dim**-0.5, + initial_state=initial_state_vk, + output_final_state=True, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + safe_gate=False, + g_is_cumsum=True, + ) + + def run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices): # SGLang/FLA GDR chunk kernels use [N, H, V, K] state layout. cuLA's # Qwen3.5 wrapper uses [N, H, K, V], so pass the transposed view here. @@ -190,12 +260,21 @@ def print_header(device: torch.device, args: argparse.Namespace, baseline_source print(f" seq lens: {args.seq_lens}") print(f" warmup/rep: {args.warmup}/{args.rep}") print(f" baselines: {baseline_sources or 'disabled/unavailable'}") + if args.cula_mode == "qk": + print(" cula: qwen35_chunk_qk_prefill_sm90 QK subpath only; baselines are full Triton GDR chunk kernels") + elif args.cula_mode == "scalar": + print(" cula: qwen35_scalar_kda_prefill full recurrence fallback") + elif args.cula_mode == "fused": + print(" cula: qwen35_fused_kda_prefill full recurrence via fused KDA CuTe core") + elif args.cula_mode == "fused-core": + print(" cula: fused KDA CuTe core only; Qwen gate/l2norm/cumsum/state prep is outside timing") print() + cula_col = f"cula_{args.cula_mode}_ms" print( - f"{'baseline':>8} {'B':>3} {'T':>7} {'layout_ms':>11} {'cula_kda_ms':>12} {'cula_total':>11} " - f"{'base_ms':>11} {'speedup':>9} {'rel_rms':>10} {'rel_max':>10}" + f"{'baseline':>8} {'B':>3} {'T':>7} {'layout_ms':>11} {cula_col:>13} {'cula_total':>11} " + f"{'base_ms':>11} {'base/cula':>10} {'rel_rms':>11} {'rel_max':>11}" ) - print("-" * 108) + print("-" * 113) def main() -> None: @@ -207,7 +286,19 @@ def main() -> None: parser.add_argument("--seed", type=int, default=0) parser.add_argument("--baseline", choices=["none", "fla", "sgl", "all"], default="sgl") parser.add_argument("--sglang-path", type=pathlib.Path, default=None) + parser.add_argument( + "--cula-mode", + choices=["qk", "scalar", "fused", "fused-core"], + default="qk", + help="cuLA path to benchmark: qk is QK subpath, scalar is old full fallback, fused is wrapper, fused-core is kernel only.", + ) parser.add_argument("--skip-accuracy", action="store_true") + parser.add_argument( + "--max-qk-elements", + type=int, + default=512 * 1024 * 1024, + help="Skip cuLA QK timings when B*48*T*T exceeds this element count.", + ) args = parser.parse_args() if not torch.cuda.is_available(): @@ -245,44 +336,88 @@ def main() -> None: def layout_fn(): return qwen35_layout_prefill(mixed_qkv_conv, a_flat, b_flat, backend="cudac") - def cula_kda_fn(): - return run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state) + qk_elements = args.batch * CONFIG.num_v_heads * seq_len * seq_len + qk_out = None + if args.cula_mode == "qk" and qk_elements <= args.max_qk_elements: + qk_out = torch.empty( + args.batch, + CONFIG.num_v_heads, + seq_len, + seq_len, + device=device, + dtype=torch.float32, + ) + fused_core_inputs = None + if args.cula_mode == "fused-core": + fused_core_inputs = prepare_cula_fused_core_inputs(q, k, a, b, A_log, dt_bias, initial_state) + + def cula_fn(): + if args.cula_mode == "scalar": + return run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state) + if args.cula_mode == "fused": + return run_cula_fused(q, k, v, a, b, A_log, dt_bias, initial_state) + if args.cula_mode == "fused-core": + return run_cula_fused_core(*fused_core_inputs[:2], v, *fused_core_inputs[2:]) + if qk_out is None: + raise RuntimeError( + f"Skipping cuLA QK: B*H*T*T={qk_elements} exceeds --max-qk-elements={args.max_qk_elements}" + ) + return run_cula_chunk_qk(q, k, qk_out) layout_ms = benchmark_cuda_fn(layout_fn, warmup=args.warmup, rep=args.rep) - cula_kda_ms = benchmark_cuda_fn(cula_kda_fn, warmup=args.warmup, rep=args.rep) - cula_total_ms = layout_ms + cula_kda_ms + cula_ms = ( + float("nan") + if args.cula_mode == "qk" and qk_out is None + else benchmark_cuda_fn(cula_fn, warmup=args.warmup, rep=args.rep) + ) + cula_total_ms = layout_ms + cula_ms if not torch.isnan(torch.tensor(cula_ms)) else float("nan") + + rel_rms = float("nan") + rel_max = float("nan") + state_cula = None + if not args.skip_accuracy: + out_cula = cula_fn() + if args.cula_mode == "qk": + qk_ref = torch.einsum("bthd,bshd->bhts", q.float(), k.float()) + torch.cuda.synchronize() + rel_rms, rel_max, _ = error_stats(qk_ref, out_cula) + del qk_ref + else: + out_cula, state_cula = out_cula + torch.cuda.synchronize() if not baselines: print( - f"{'none':>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_kda_ms:12.4f} {cula_total_ms:11.4f} " - f"{float('nan'):11.4f} {float('nan'):9.3f} {float('nan'):10.3e} {float('nan'):10.3e}" + f"{'none':>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_ms:13.4f} {cula_total_ms:11.4f} " + f"{float('nan'):11.4f} {float('nan'):10.3f} {rel_rms:11.3e} {rel_max:11.3e}" ) for baseline_name, chunk_gdr in baselines.items(): def baseline_fn(): return run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices) - rel_rms = float("nan") - rel_max = float("nan") - if not args.skip_accuracy: - out_cula, state_cula = cula_kda_fn() + row_rel_rms = rel_rms + row_rel_max = rel_max + if args.cula_mode in ("scalar", "fused", "fused-core") and not args.skip_accuracy: + if state_cula is None: + out_cula, state_cula = cula_fn() out_base, state_base = _normalize_chunk_result(baseline_fn()) state_base = _state_to_cula_layout(state_base) torch.cuda.synchronize() - rel_rms, rel_max, _ = error_stats(out_base, out_cula) + row_rel_rms, row_rel_max, _ = error_stats(out_base, out_cula) if state_base is not None and tuple(state_base.shape) == tuple(state_cula.shape): rel_rms_s, rel_max_s, _ = error_stats(state_base, state_cula) - rel_rms = max(rel_rms, rel_rms_s) - rel_max = max(rel_max, rel_max_s) + row_rel_rms = max(row_rel_rms, rel_rms_s) + row_rel_max = max(row_rel_max, rel_max_s) base_ms = benchmark_cuda_fn(baseline_fn, warmup=args.warmup, rep=args.rep) - speedup = base_ms / cula_kda_ms if cula_kda_ms > 0 else float("inf") + speedup = base_ms / cula_ms if cula_ms > 0 else float("nan") print( - f"{baseline_name:>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_kda_ms:12.4f} {cula_total_ms:11.4f} " - f"{base_ms:11.4f} {speedup:9.3f} {rel_rms:10.3e} {rel_max:10.3e}" + f"{baseline_name:>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_ms:13.4f} {cula_total_ms:11.4f} " + f"{base_ms:11.4f} {speedup:10.3f} {row_rel_rms:11.3e} {row_rel_max:11.3e}" ) - del q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, initial_state_indices, mixed_qkv_conv, a_flat, b_flat + del q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, initial_state_indices, mixed_qkv_conv, a_flat, b_flat, qk_out, fused_core_inputs torch.cuda.empty_cache() diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index 2e93cf0a..3f741945 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -211,6 +211,11 @@ qwen35_layout_prefill( cula::qwen35::prefill::run_qwen35_layout_prefill(params); } +void +qwen35_chunk_qk_prefill_sm90(at::Tensor q, at::Tensor k, at::Tensor out) { + cula::qwen35::prefill::sm90::qwen35_chunk_qk_prefill_sm90(q, k, out); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "cuLA"; #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) @@ -226,4 +231,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("qwen35_layout_scalar_kda_decode", &qwen35_layout_scalar_kda_decode); m.def("qwen35_layout_prefill", &qwen35_layout_prefill); m.def("qwen35_scalar_kda_prefill", &qwen35_scalar_kda_prefill); + m.def("qwen35_chunk_qk_prefill_sm90", &qwen35_chunk_qk_prefill_sm90); } diff --git a/csrc/qwen35/prefill/qwen35_prefill_common.cuh b/csrc/qwen35/prefill/qwen35_prefill_common.cuh index 0aac63b2..6c7c6bf7 100644 --- a/csrc/qwen35/prefill/qwen35_prefill_common.cuh +++ b/csrc/qwen35/prefill/qwen35_prefill_common.cuh @@ -58,3 +58,9 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params); void run_qwen35_layout_prefill(LayoutPrefillParams& params); } // namespace cula::qwen35::prefill + +namespace cula::qwen35::prefill::sm90 { + +void qwen35_chunk_qk_prefill_sm90(at::Tensor q, at::Tensor k, at::Tensor out); + +} // namespace cula::qwen35::prefill::sm90 diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp index f0aba77f..db135475 100644 --- a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp @@ -27,7 +27,11 @@ template struct Qwen35ScalarKdaPrefillKernel { static constexpr int kThreads = 128; static constexpr int kHeadDim = kHeadDimQK; - static constexpr int kVTile = 8; + // Keep the scalar CUDA fallback at one V row per CTA for correctness while + // the SM90 chunk/TMA path is being wired in. The previous multi-row V tile + // version exposed a correctness bug with non-zero initial_state; the chunk + // path should own the next parallelization step. + static constexpr int kVTile = 1; static constexpr int kNumVTiles = kHeadDimV / kVTile; static_assert(kHeadDimQK == 128); diff --git a/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu new file mode 100644 index 00000000..5af11afe --- /dev/null +++ b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu @@ -0,0 +1,180 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qwen35_chunk_prefill_traits_sm90.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cula::qwen35::prefill::sm90 { + +namespace { + +using DefaultTraits = Qwen35ChunkPrefillSm90DefaultTraits; + +static_assert(DefaultTraits::kBlockT == 64); +static_assert(DefaultTraits::kBlockV == 64); +static_assert(DefaultTraits::kStages == 2); +static_assert(size(typename DefaultTraits::TiledMmaQK{}) == 128); +static_assert(size(typename DefaultTraits::TiledMmaOV{}) == 128); +static_assert(cosize(typename DefaultTraits::SmemLayoutQ{}) > 0); +static_assert(cosize(typename DefaultTraits::SmemLayoutK{}) > 0); + +void check_cutlass_status(cutlass::Status status, const char* what) { + TORCH_CHECK(status == cutlass::Status::kSuccess, what, " failed with CUTLASS status ", static_cast(status)); +} + +template +void run_qwen35_chunk_qk_prefill_sm90_impl(const at::Tensor& q, const at::Tensor& k, const at::Tensor& out) { + using ElementA = cutlass::bfloat16_t; + using ElementB = cutlass::bfloat16_t; + using ElementC = float; + using ElementD = float; + using ElementAccumulator = float; + using ElementCompute = float; + + using LayoutA = cute::tuple; + using LayoutB = cute::tuple; + using LayoutC = cute::tuple; + using LayoutD = LayoutC; + + constexpr int kAlignmentA = 16 / sizeof(ElementA); + constexpr int kAlignmentB = 16 / sizeof(ElementB); + constexpr int kAlignmentC = 16 / sizeof(ElementC); + constexpr int kAlignmentD = 16 / sizeof(ElementD); + + using OperatorClass = cutlass::arch::OpClassTensorOp; + using TileShape = cute::Shape; + using ClusterShape = cute::Shape; +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) + using ArchTag = cutlass::arch::Sm100; + using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; + using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; +#else + using ArchTag = cutlass::arch::Sm90; + using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized; + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized; +#endif +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) + using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; +#else + using EpilogueTileType = decltype(cute::take<0, 2>(TileShape{})); +#endif + using FusionOperation = + typename cutlass::epilogue::fusion::LinearCombination; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + TileShape, + ClusterShape, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC, + LayoutC, + kAlignmentC, + ElementD, + LayoutD, + kAlignmentD, + EpilogueSchedule, + FusionOperation>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + TileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal, CollectiveMainloop, CollectiveEpilogue>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t HV = q.size(2); + constexpr int K = kHeadDimQK; + const int64_t L = B * HV; + + LayoutA stride_A{HV * K, cute::_1{}, K}; + LayoutB stride_B{HV * K, cute::_1{}, K}; + LayoutC stride_C{T, cute::_1{}, T * T}; + + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {static_cast(T), static_cast(T), K, static_cast(L)}, + { + reinterpret_cast(q.data_ptr()), + stride_A, + reinterpret_cast(k.data_ptr()), + stride_B, + }, + { + {1.0f, 0.0f}, + out.data_ptr(), + stride_C, + out.data_ptr(), + stride_C, + }, + }; + + Gemm gemm; + const size_t workspace_size = Gemm::get_workspace_size(arguments); + at::Tensor workspace = at::empty({static_cast(workspace_size)}, q.options().dtype(at::kByte)); + check_cutlass_status(gemm.can_implement(arguments), "qwen35_chunk_qk_prefill_sm90 can_implement"); + check_cutlass_status(gemm.initialize(arguments, workspace.data_ptr(), at::cuda::getCurrentCUDAStream(q.device().index())), "qwen35_chunk_qk_prefill_sm90 initialize"); + check_cutlass_status(gemm.run(at::cuda::getCurrentCUDAStream(q.device().index())), "qwen35_chunk_qk_prefill_sm90 run"); +} + +} // namespace + +void qwen35_chunk_qk_prefill_sm90(at::Tensor q, at::Tensor k, at::Tensor out) { + TORCH_CHECK(q.is_cuda(), "q must be CUDA"); + TORCH_CHECK(k.is_cuda(), "k must be CUDA"); + TORCH_CHECK(out.is_cuda(), "out must be CUDA"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16, "q must be bfloat16"); + TORCH_CHECK(k.scalar_type() == at::kBFloat16, "k must be bfloat16"); + TORCH_CHECK(out.scalar_type() == at::kFloat, "out must be float32"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B,T,48,128]"); + TORCH_CHECK(k.is_contiguous(), "k must be contiguous [B,T,48,128]"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B,48,T,T]"); + TORCH_CHECK(q.dim() == 4, "q must be [B,T,48,128]"); + TORCH_CHECK(k.sizes() == q.sizes(), "k must match q"); + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t HV = q.size(2); + TORCH_CHECK(HV == kNumVHeads, "expected HV=48"); + TORCH_CHECK(q.size(3) == kHeadDimQK, "expected D=128"); + TORCH_CHECK(out.sizes() == at::IntArrayRef({B, HV, T, T}), "out must be [B,48,T,T]"); + + const at::cuda::OptionalCUDAGuard device_guard(q.device()); + run_qwen35_chunk_qk_prefill_sm90_impl(q, k, out); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace cula::qwen35::prefill::sm90 diff --git a/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_traits_sm90.hpp b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_traits_sm90.hpp new file mode 100644 index 00000000..c39e926e --- /dev/null +++ b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_traits_sm90.hpp @@ -0,0 +1,112 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "qwen35/prefill/qwen35_prefill_common.cuh" + +#include +#include +#include +#include +#include +#include + +namespace cula::qwen35::prefill::sm90 { + +using namespace cute; + +// First SM90 chunk shape for Qwen3.5 prefill. +// +// This intentionally only describes the TMA/WGMMA tiles. The full chunk +// algorithm still needs a local chunk recurrence and inter-chunk state scan; +// those should be built on top of these traits instead of extending the scalar +// fallback kernel. +template +struct Qwen35ChunkPrefillSm90Traits { + static constexpr int kBlockT = kBlockT_; + static constexpr int kBlockV = kBlockV_; + static constexpr int kStages = kStages_; + + static_assert(kBlockT == 64 || kBlockT == 128, "GMMA chunk tiles expect BT=64 or BT=128."); + static_assert(kBlockV == 64 || kBlockV == 128, "V chunk tiles expect BV=64 or BV=128."); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + + using Element = cutlass::bfloat16_t; + using Accumulator = float; + static constexpr int kAlignment = 16 / sizeof(Element); + + using ClusterShape = Shape<_1, _1, _1>; + using StageCount = cutlass::gemm::collective::StageCount; + + // q/k/v are materialized by qwen35_layout_prefill as contiguous + // [total_tokens, 48, 128]. The TMA tensor view below exposes them as + // (token, dim, head), with dynamic strides: + // token stride = 48 * 128 + // dim stride = 1 + // head stride = 128 + using GmemStrideTDH = cute::tuple; + + using TileShapeQK = decltype(make_shape(Int{}, Int{}, Int{})); + using TileShapeOV = decltype(make_shape(Int{}, Int{}, Int{})); + + // Q @ K^T => [BT, BT]. CollectiveBuilder selects GMMA and TMA-compatible + // shared-memory layouts for SM90. + using CollectiveQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + GmemStrideTDH, + kAlignment, + Element, + GmemStrideTDH, + kAlignment, + Accumulator, + TileShapeQK, + ClusterShape, + StageCount, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaQK = typename CollectiveQK::TiledMma; + using SmemLayoutQ = typename CollectiveQK::SmemLayoutA; + using SmemLayoutK = typename CollectiveQK::SmemLayoutB; + using TmaQ = typename CollectiveQK::Params::TMA_A; + using TmaK = typename CollectiveQK::Params::TMA_B; + + // Q @ state / local_value => [BT, BV]. This is the second core WGMMA shape + // needed once chunk-local state summaries are available. + using CollectiveOV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm90, + cutlass::arch::OpClassTensorOp, + Element, + GmemStrideTDH, + kAlignment, + Element, + GmemStrideTDH, + kAlignment, + Accumulator, + TileShapeOV, + ClusterShape, + StageCount, + cutlass::gemm::KernelTmaWarpSpecialized>::CollectiveOp; + + using TiledMmaOV = typename CollectiveOV::TiledMma; + using SmemLayoutOV_A = typename CollectiveOV::SmemLayoutA; + using SmemLayoutOV_B = typename CollectiveOV::SmemLayoutB; +}; + +using Qwen35ChunkPrefillSm90DefaultTraits = Qwen35ChunkPrefillSm90Traits<64, 64, 2>; + +} // namespace cula::qwen35::prefill::sm90 diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index ee1a2bb9..98f5a14d 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -13,9 +13,19 @@ # limitations under the License. from cula.kda.blackwell_fused_fwd import flash_kda_prefill as kda_prefill_blackwell -from cula.kda.chunk import chunk_kda -from cula.kda.hopper_fused_fwd import cula_kda_prefill as kda_prefill_hopper -from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode +try: + from cula.kda.chunk import chunk_kda +except Exception: # pragma: no cover - optional FLA dependency + chunk_kda = None +try: + from cula.kda.hopper_fused_fwd import cula_kda_prefill as kda_prefill_hopper +except Exception: # pragma: no cover - optional FLA/Hopper dependency + kda_prefill_hopper = None +try: + from cula.ops.kda_decode import fused_sigmoid_gating_delta_rule_update, kda_decode +except Exception: # pragma: no cover - optional CUDA/CuTe dependency + fused_sigmoid_gating_delta_rule_update = None + kda_decode = None __all__ = [ "chunk_kda", diff --git a/cula/kda/blackwell_fused_fwd.py b/cula/kda/blackwell_fused_fwd.py index c7dec95c..5ac821ba 100644 --- a/cula/kda/blackwell_fused_fwd.py +++ b/cula/kda/blackwell_fused_fwd.py @@ -17,6 +17,7 @@ import warnings import torch +import torch.nn.functional as F sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) @@ -24,13 +25,58 @@ import cutlass.cute as cute import cutlass.torch as cutlass_torch from cutlass.cute.runtime import from_dlpack -from fla.modules.l2norm import l2norm_fwd # from fla.ops.kda.chunk_inter import chunk_kda_bwd_dqkwg -from fla.ops.kda.gate import kda_gate_fwd -from fla.ops.utils import chunk_local_cumsum -from fla.ops.utils.constant import RCP_LN2 -from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +try: + from fla.modules.l2norm import l2norm_fwd + from fla.ops.kda.gate import kda_gate_fwd + from fla.ops.utils import chunk_local_cumsum + from fla.ops.utils.constant import RCP_LN2 + from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +except ImportError: + RCP_LN2 = 1.4426950408889634 + + def input_guard(fn): + return fn + + def autocast_custom_fwd(fn): + return fn + + def autocast_custom_bwd(fn): + return fn + + def l2norm_fwd(x: torch.Tensor): + rstd = torch.rsqrt(x.float().square().sum(dim=-1, keepdim=True).clamp_min(1.0e-12)) + return (x.float() * rstd).to(x.dtype), rstd + + def kda_gate_fwd(*args, **kwargs): + raise ImportError("fla is required for use_gate_in_kernel=True in blackwell_fused_fwd") + + def chunk_local_cumsum( + g: torch.Tensor, + chunk_size: int, + scale: float = 1.0, + cu_seqlens: torch.Tensor | None = None, + chunk_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + if chunk_indices is not None: + raise ImportError("fla is required for chunk_indices support in blackwell_fused_fwd") + if cu_seqlens is not None: + if g.shape[0] != 1: + raise ValueError("cu_seqlens mode expects flattened g with batch size 1") + out = torch.empty_like(g.float()) + for seq_idx in range(cu_seqlens.numel() - 1): + start = int(cu_seqlens[seq_idx].item()) + end = int(cu_seqlens[seq_idx + 1].item()) + for chunk_start in range(start, end, chunk_size): + chunk_end = min(chunk_start + chunk_size, end) + out[:, chunk_start:chunk_end] = g[:, chunk_start:chunk_end].float().cumsum(dim=1) * scale + return out + chunks = [] + for chunk_start in range(0, g.shape[1], chunk_size): + chunk = g[:, chunk_start : chunk_start + chunk_size].float().cumsum(dim=1) * scale + chunks.append(chunk) + return torch.cat(chunks, dim=1).contiguous() from cula.ops.kda_fully_fused_sm100_wip import KDAChunkwise from cula.utils import USE_FAST_MATH, assert_blackwell @@ -64,6 +110,7 @@ def forward( use_gate_in_kernel: bool = False, safe_gate: bool = False, lower_bound: float | None = None, + g_is_cumsum: bool = False, cu_seqlens: torch.IntTensor | None = None, chunk_indices: torch.IntTensor | None = None, ): @@ -106,7 +153,7 @@ def forward( A_log=A_log, dt_bias=dt_bias, ) - if not (safe_gate and use_gate_in_kernel): + if not g_is_cumsum and not (safe_gate and use_gate_in_kernel): g = chunk_local_cumsum( g=g, chunk_size=chunk_size, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices ) @@ -267,6 +314,7 @@ def flash_kda_prefill( use_gate_in_kernel: bool = False, safe_gate: bool = False, lower_bound: float | None = None, + g_is_cumsum: bool = False, cu_seqlens: torch.IntTensor | None = None, chunk_indices: torch.IntTensor | None = None, **kwargs, @@ -326,6 +374,7 @@ def flash_kda_prefill( use_gate_in_kernel, safe_gate, lower_bound, + g_is_cumsum, cu_seqlens, chunk_indices, ) diff --git a/cula/ops/kda_fully_fused_sm100_wip.py b/cula/ops/kda_fully_fused_sm100_wip.py index ab09f0b2..ec3d5fe9 100644 --- a/cula/ops/kda_fully_fused_sm100_wip.py +++ b/cula/ops/kda_fully_fused_sm100_wip.py @@ -68,7 +68,12 @@ from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Int32, Int64 -from fla.modules.l2norm import l2norm_fwd +try: + from fla.modules.l2norm import l2norm_fwd +except ImportError: + def l2norm_fwd(x: torch.Tensor): + rstd = torch.rsqrt(x.float().square().sum(dim=-1, keepdim=True).clamp_min(1.0e-12)) + return (x.float() * rstd).to(x.dtype), rstd from cula.utils import assert_blackwell @@ -2993,13 +2998,16 @@ def index_transform_half(index_q, index_k): # ------------------------------------------------------------ # NOTE: Save exp(g) of last VALID row to rG_last for state update in next chunk # For full chunks, directly use C-1; only loop for partial chunks (varlen only) + rG_last = cutlass.Float32(1.0) if cutlass.const_expr(self.is_varlen): if valid_len_chunk < C: - rG_last = exp_g[valid_len_chunk - 1] + for _zr in cutlass.range(0, Constant.C, unroll_full=True): + if _zr == valid_len_chunk - 1: + rG_last = cute.exp2(tRS_rG[0, _zr, 0], fastmath=self.use_fast_math) else: - rG_last = exp_g[Constant.C - 1] + rG_last = cute.exp2(tRS_rG[0, Constant.C - 1, 0], fastmath=self.use_fast_math) else: - rG_last = exp_g[Constant.C - 1] + rG_last = cute.exp2(tRS_rG[0, Constant.C - 1, 0], fastmath=self.use_fast_math) # NOTE: each thread save one element sG_last[local_tidx, g_stage_idx] = rG_last diff --git a/cula/ops/qwen35_fused_kda_prefill.py b/cula/ops/qwen35_fused_kda_prefill.py new file mode 100644 index 00000000..6691fdd8 --- /dev/null +++ b/cula/ops/qwen35_fused_kda_prefill.py @@ -0,0 +1,132 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen3.5 adapter for the generic fused KDA prefill core.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def _resolve_fused_kda_prefill(device: torch.device | str | int | None = None): + try: + from cula.utils import get_kda_fused_fwd + except Exception as exc: # pragma: no cover - depends on optional runtime deps + raise RuntimeError(f"Cannot import fused KDA selector: {exc}") from exc + + try: + return get_kda_fused_fwd(device) + except Exception as exc: + raise RuntimeError(f"Cannot resolve fused KDA prefill for device={device}: {exc}") from exc + + +def has_qwen35_fused_kda_prefill(device: torch.device | str | int | None = None) -> bool: + try: + _resolve_fused_kda_prefill(device) + except Exception: + return False + return True + + +def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +) -> tuple[int, int, int, int, torch.Tensor, torch.Tensor]: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError(f"q/k/v must be 4D [B,T,HV,D], got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + if q.shape != k.shape or q.shape != v.shape: + raise ValueError(f"q/k/v must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + B, T, HV, K = q.shape + if K != 128: + raise ValueError(f"Qwen3.5 fused prefill expects head dim 128, got {K}") + if a.ndim == 2: + a = a.unsqueeze(0) + if b.ndim == 2: + b = b.unsqueeze(0) + if a.shape != (B, T, HV) or b.shape != (B, T, HV): + raise ValueError(f"a/b must be [B,T,HV], got a={tuple(a.shape)} b={tuple(b.shape)} expected={(B, T, HV)}") + if A_log.shape != (HV,) or dt_bias.shape != (HV,): + raise ValueError(f"A_log/dt_bias must be [HV], got A_log={tuple(A_log.shape)} dt_bias={tuple(dt_bias.shape)}") + if cu_seqlens is not None: + if B != 1: + raise ValueError("cu_seqlens mode expects flattened q/k/v with batch size 1") + if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + if initial_state is not None and initial_state.shape != (state_count, HV, K, K): + raise ValueError(f"initial_state must be [{state_count},{HV},128,128], got {tuple(initial_state.shape)}") + return B, T, HV, K, a, b + + +def qwen35_fused_kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + output_final_state: bool = True, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run Qwen3.5 scalar-gated KDA prefill through the fused CuTe KDA core. + + Qwen uses a scalar gate per token/head. The generic KDA fused core expects + a vector gate, so this adapter broadcasts the scalar log-gate over D=128. + State is exposed in Qwen layout [N, HV, K, V]. The fused core consumes the + transposed initial-state layout, but returns final state in Qwen layout. + """ + + if not q.is_cuda: + raise RuntimeError("qwen35_fused_kda_prefill requires CUDA tensors.") + B, T, HV, K, a, b = _validate_inputs(q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens) + fused_kda_prefill = _resolve_fused_kda_prefill(q.device) + + log_gate_scalar = -torch.exp(A_log.float()).view(1, 1, HV, 1) * F.softplus( + a.float().unsqueeze(-1) + dt_bias.float().view(1, 1, HV, 1) + ) + log_gate = log_gate_scalar.expand(B, T, HV, K).contiguous() + beta = torch.sigmoid(b.float()).contiguous() + + initial_state_vk = None + if initial_state is not None: + initial_state_vk = initial_state.float().transpose(-1, -2).contiguous() + + out, final_state_vk = fused_kda_prefill( + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + g=log_gate, + beta=beta, + scale=K**-0.5, + initial_state=initial_state_vk, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=False, + safe_gate=False, + lower_bound=None, + cu_seqlens=cu_seqlens, + ) + final_state = None if final_state_vk is None else final_state_vk.contiguous() + return out, final_state diff --git a/cula/utils.py b/cula/utils.py index 8b8e0ab1..ed020f73 100644 --- a/cula/utils.py +++ b/cula/utils.py @@ -94,11 +94,11 @@ def get_kda_fused_fwd(device: torch.device | str | int | None = None) -> Callabl """ major, minor = get_device_sm_version(device) if major == 10 and minor in (0, 3): - from cula.kda import kda_prefill_blackwell + from cula.kda.blackwell_fused_fwd import flash_kda_prefill as kda_prefill_blackwell return kda_prefill_blackwell elif major == 9 and minor == 0: - from cula.kda import kda_prefill_hopper + from cula.kda.hopper_fused_fwd import cula_kda_prefill as kda_prefill_hopper return kda_prefill_hopper else: diff --git a/setup.py b/setup.py index 8ee1bc55..11d7f654 100644 --- a/setup.py +++ b/setup.py @@ -152,6 +152,7 @@ def get_nvcc_thread_args(): "csrc/qwen35/decode/qwen35_scalar_kda_decode.cu", "csrc/qwen35/prefill/qwen35_layout_prefill.cu", "csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu", + "csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu", ] if not DISABLE_SM100 or not DISABLE_SM103: cuda_sources.extend( diff --git a/tests/test_qwen35_prefill.py b/tests/test_qwen35_prefill.py index a1e0b640..d598b71a 100644 --- a/tests/test_qwen35_prefill.py +++ b/tests/test_qwen35_prefill.py @@ -21,11 +21,17 @@ sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) from cula.ops.qwen35_conv1d_prefill import qwen35_conv1d_prefill +from cula.ops.qwen35_fused_kda_prefill import has_qwen35_fused_kda_prefill, qwen35_fused_kda_prefill from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill, qwen35_layout_prefill_reference from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill from cula.qwen35.common import Qwen35LinearAttentionConfig from cula.qwen35.runtime import qwen35_linear_attention_prefill +try: + import cula.cudac as cula_cuda +except ImportError: + cula_cuda = None + def _manual_scalar_prefill(q, k, v, a, b, A_log, dt_bias, initial_state=None, cu_seqlens=None): B, T, HV, K = q.shape @@ -118,6 +124,121 @@ def test_qwen35_scalar_kda_prefill_varlen_reference_matches_manual(): torch.testing.assert_close(state, state_ref, atol=1e-4, rtol=1e-4) +def test_qwen35_scalar_kda_prefill_cuda_matches_reference(): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + import pytest + + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(10) + device = torch.device("cuda") + B, T, HV, K = 1, 8, 48, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + +def test_qwen35_chunk_qk_prefill_sm90_matches_torch(): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): + import pytest + + pytest.skip("qwen35_chunk_qk_prefill_sm90 CUDA extension is not available") + + torch.manual_seed(11) + device = torch.device("cuda") + B, T, HV, K = 1, 64, 48, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + out = torch.empty(B, HV, T, T, device=device, dtype=torch.float32) + + cula_cuda.qwen35_chunk_qk_prefill_sm90(q.contiguous(), k.contiguous(), out) + torch.cuda.synchronize() + + ref = torch.einsum("bthd,bshd->bhts", q.float(), k.float()) + torch.testing.assert_close(out, ref, atol=2e-1, rtol=2e-2) + + +def test_qwen35_fused_kda_prefill_matches_reference(): + if not torch.cuda.is_available(): + import pytest + + pytest.skip("CUDA is not available") + if not has_qwen35_fused_kda_prefill(torch.device("cuda")): + import pytest + + pytest.skip("Qwen3.5 fused KDA prefill backend is not available") + + torch.manual_seed(12) + device = torch.device("cuda") + B, T, HV, K = 1, 64, 48, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + out, state = qwen35_fused_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state, state_ref, atol=3e-2, rtol=3e-2) + + def test_qwen35_conv1d_prefill_flattened_state(): x = torch.arange(5 * 3, dtype=torch.bfloat16).reshape(5, 3) weight = torch.ones(3, 4, dtype=torch.bfloat16) From 9c26912339ad6ac46b9f45836b9d4ecfd4aee640 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Thu, 11 Jun 2026 13:16:48 +0000 Subject: [PATCH 10/35] Support Qwen3.5 local TP head configs --- benchmarks/bench_qwen35_decode.py | 146 +++++++++------ benchmarks/bench_qwen35_prefill.py | 76 +++++--- csrc/qwen35/decode/qwen35_conv1d_decode.cu | 44 +++-- csrc/qwen35/decode/qwen35_decode_common.cuh | 76 +++++--- csrc/qwen35/decode/qwen35_layout_decode.cu | 90 ++++++--- csrc/qwen35/decode/qwen35_layout_kernel.hpp | 33 ++-- .../qwen35/decode/qwen35_scalar_kda_decode.cu | 174 ++++++++++++------ .../decode/qwen35_scalar_kda_kernel.hpp | 95 +++++----- csrc/qwen35/prefill/qwen35_layout_prefill.cu | 75 +++++++- .../prefill/qwen35_layout_prefill_kernel.hpp | 39 ++-- csrc/qwen35/prefill/qwen35_prefill_common.cuh | 36 ++-- .../prefill/qwen35_scalar_kda_prefill.cu | 97 +++++++++- .../qwen35_scalar_kda_prefill_kernel.hpp | 34 ++-- .../prefill/sm90/qwen35_chunk_prefill_sm90.cu | 12 +- cula/ops/qwen35_layout_decode.py | 1 + cula/ops/qwen35_layout_prefill.py | 3 +- cula/ops/qwen35_scalar_kda_prefill.py | 4 +- tests/test_qwen35_decode.py | 62 ++++++- tests/test_qwen35_prefill.py | 90 +++++++++ 19 files changed, 832 insertions(+), 355 deletions(-) diff --git a/benchmarks/bench_qwen35_decode.py b/benchmarks/bench_qwen35_decode.py index 932f757b..7cd9f858 100755 --- a/benchmarks/bench_qwen35_decode.py +++ b/benchmarks/bench_qwen35_decode.py @@ -51,6 +51,7 @@ import cula.cudac as cula_cuda from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as triton_fused_sigmoid_update from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as CONFIG +from cula.qwen35.common import Qwen35LinearAttentionConfig from cula.qwen35.runtime import qwen35_linear_attention_decode SGLANG_CORE_MODULES = [ @@ -134,71 +135,86 @@ def benchmark_accel_fn( return statistics.mean(iqr) -def make_full_inputs(tokens: int, device: torch.device, seed: int): +def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: + if tp_size not in (1, 2, 4, 8): + raise ValueError(f"tp_size must be one of 1, 2, 4, 8, got {tp_size}") + return Qwen35LinearAttentionConfig( + hidden_size=CONFIG.hidden_size // tp_size, + conv_kernel_size=CONFIG.conv_kernel_size, + num_k_heads=CONFIG.num_k_heads // tp_size, + num_v_heads=CONFIG.num_v_heads // tp_size, + head_k_dim=CONFIG.head_k_dim, + head_v_dim=CONFIG.head_v_dim, + qkv_dtype=CONFIG.qkv_dtype, + state_dtype=CONFIG.state_dtype, + ) + + +def make_full_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): torch.manual_seed(seed) pool_size = max(tokens, 1) - mixed_qkv = torch.randn(tokens, CONFIG.conv_dim, device=device, dtype=CONFIG.qkv_dtype) - a = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - b = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - conv_weight = torch.randn(CONFIG.conv_dim, CONFIG.conv_kernel_size, device=device, dtype=CONFIG.qkv_dtype) + mixed_qkv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + conv_weight = torch.randn(config.conv_dim, config.conv_kernel_size, device=device, dtype=config.qkv_dtype) conv_state = torch.randn( tokens, - CONFIG.conv_dim, - CONFIG.conv_kernel_size, + config.conv_dim, + config.conv_kernel_size, device=device, - dtype=CONFIG.qkv_dtype, + dtype=config.qkv_dtype, ) recurrent_state = torch.randn( pool_size, - CONFIG.num_v_heads, - CONFIG.head_k_dim, - CONFIG.head_v_dim, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, device=device, - dtype=CONFIG.state_dtype, + dtype=config.state_dtype, ) * 0.01 - A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 state_indices = torch.arange(tokens, device=device, dtype=torch.int32) return mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices -def make_fused_layout_kda_inputs(tokens: int, device: torch.device, seed: int): +def make_fused_layout_kda_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): torch.manual_seed(seed) - mixed_qkv_conv = torch.randn(tokens, CONFIG.conv_dim, device=device, dtype=CONFIG.qkv_dtype) - a = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - b = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + mixed_qkv_conv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 state = torch.randn( tokens, - CONFIG.num_v_heads, - CONFIG.head_k_dim, - CONFIG.head_v_dim, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, device=device, - dtype=CONFIG.state_dtype, + dtype=config.state_dtype, ) * 0.01 state_work = torch.empty_like(state) state_indices = torch.arange(tokens, device=device, dtype=torch.int32) - out = torch.empty(tokens, CONFIG.num_v_heads, CONFIG.head_v_dim, device=device, dtype=CONFIG.qkv_dtype) + out = torch.empty(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) return mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out -def make_core_inputs(tokens: int, device: torch.device, seed: int): +def make_core_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): torch.manual_seed(seed) - q = torch.randn(tokens, CONFIG.num_v_heads, CONFIG.head_k_dim, device=device, dtype=CONFIG.qkv_dtype) - k = torch.randn(tokens, CONFIG.num_v_heads, CONFIG.head_k_dim, device=device, dtype=CONFIG.qkv_dtype) - v = torch.randn(tokens, CONFIG.num_v_heads, CONFIG.head_v_dim, device=device, dtype=CONFIG.qkv_dtype) - a = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - b = torch.randn(tokens, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + q = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + k = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + v = torch.randn(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 state = torch.randn( tokens, - CONFIG.num_v_heads, - CONFIG.head_k_dim, - CONFIG.head_v_dim, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, device=device, - dtype=CONFIG.state_dtype, + dtype=config.state_dtype, ) * 0.01 state_indices = torch.arange(tokens, device=device, dtype=torch.int32) out = torch.empty_like(v) @@ -291,8 +307,8 @@ def call_with_supported_kwargs(fn: Callable, **kwargs): return fn(**filtered) -def bench_native_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: - q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_core_inputs(tokens, device, seed) +def bench_native_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: + q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_core_inputs(tokens, device, seed, config) def setup() -> None: state_work.copy_(state) @@ -314,8 +330,8 @@ def run() -> None: return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) -def bench_triton_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: - q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, _ = make_core_inputs(tokens, device, seed) +def bench_triton_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: + q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, _ = make_core_inputs(tokens, device, seed, config) q_4d = q.unsqueeze(1).contiguous() k_4d = k.unsqueeze(1).contiguous() v_4d = v.unsqueeze(1).contiguous() @@ -338,7 +354,7 @@ def run() -> None: b=b_3d, initial_state_source=state_work, initial_state_indices=state_indices, - scale=CONFIG.head_k_dim**-0.5, + scale=config.head_k_dim**-0.5, use_qk_l2norm_in_kernel=True, cu_seqlens=None, is_kda=False, @@ -347,9 +363,9 @@ def run() -> None: return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) -def bench_fused_layout_kda(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: +def bench_fused_layout_kda(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_fused_layout_kda_inputs( - tokens, device, seed + tokens, device, seed, config ) def setup() -> None: @@ -377,8 +393,9 @@ def bench_sglang_core( rep: int, seed: int, sglang_fused_update: Callable, + config: Qwen35LinearAttentionConfig, ) -> float: - q, k, v, a, b, A_log, dt_bias, state, _, state_indices, _ = make_core_inputs(tokens, device, seed) + q, k, v, a, b, A_log, dt_bias, state, _, state_indices, _ = make_core_inputs(tokens, device, seed, config) q_4d = q.unsqueeze(1).contiguous() k_4d = k.unsqueeze(1).contiguous() v_4d = v.unsqueeze(1).contiguous() @@ -404,7 +421,7 @@ def run() -> None: b=b_3d, initial_state_source=state_vk_work, initial_state_indices=state_indices, - scale=CONFIG.head_k_dim**-0.5, + scale=config.head_k_dim**-0.5, use_qk_l2norm_in_kernel=True, cu_seqlens=None, is_kda=False, @@ -420,11 +437,12 @@ def bench_sglang_packed_layout_kda( rep: int, seed: int, sglang_packed_decode: Callable, + config: Qwen35LinearAttentionConfig, ) -> float: - mixed_qkv_conv, a, b, A_log, dt_bias, state, _, state_indices, _ = make_fused_layout_kda_inputs(tokens, device, seed) + mixed_qkv_conv, a, b, A_log, dt_bias, state, _, state_indices, _ = make_fused_layout_kda_inputs(tokens, device, seed, config) state_vk = state.transpose(-1, -2).contiguous() state_vk_work = torch.empty_like(state_vk) - out = torch.empty(tokens, 1, CONFIG.num_v_heads, CONFIG.head_v_dim, device=device, dtype=CONFIG.qkv_dtype) + out = torch.empty(tokens, 1, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) def setup() -> None: state_vk_work.copy_(state_vk) @@ -437,7 +455,7 @@ def run() -> None: b=b, A_log=A_log, dt_bias=dt_bias, - scale=CONFIG.head_k_dim**-0.5, + scale=config.head_k_dim**-0.5, initial_state=state_vk_work, out=out, ssm_state_indices=state_indices, @@ -447,8 +465,8 @@ def run() -> None: return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) -def bench_full(tokens: int, device: torch.device, warmup: int, rep: int, seed: int) -> float: - inputs = make_full_inputs(tokens, device, seed) +def bench_full(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: + inputs = make_full_inputs(tokens, device, seed, config) mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = inputs conv_state_work = torch.empty_like(conv_state) recurrent_state_work = torch.empty_like(recurrent_state) @@ -465,6 +483,7 @@ def run() -> None: conv_weight, A_log, dt_bias, + config=config, conv_state=conv_state_work, recurrent_state=recurrent_state_work, state_indices=state_indices, @@ -481,6 +500,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--rep", type=int, default=100) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--scope", choices=["core", "fused", "full", "both"], default="both") + parser.add_argument("--tp-size", type=int, choices=[1, 2, 4, 8], default=1) parser.add_argument("--skip-triton", action="store_true", help="Skip the vendored Triton core timing.") parser.add_argument("--skip-sglang", action="store_true", help="Do not try the SGLang kernel provider.") parser.add_argument("--require-sglang", action="store_true", help="Fail if the SGLang kernel provider is unavailable.") @@ -491,6 +511,7 @@ def parse_args() -> argparse.Namespace: def main() -> int: args = parse_args() + config = local_config_from_tp_size(args.tp_size) device = accelerator_device() rows: list[dict[str, object]] = [] sglang_fused_update = None @@ -504,7 +525,10 @@ def main() -> int: raise RuntimeError("SGLang core and packed decode providers must both be available.") print(f"device={device} name={accelerator_name(device)} torch={torch.__version__}") - print(f"qwen35: HV={CONFIG.num_v_heads} K={CONFIG.head_k_dim} V={CONFIG.head_v_dim} conv_dim={CONFIG.conv_dim}") + print( + f"qwen35: tp={args.tp_size} local_HK={config.num_k_heads} local_HV={config.num_v_heads} " + f"K={config.head_k_dim} V={config.head_v_dim} conv_dim={config.conv_dim}" + ) print(f"sglang_core_provider={sglang_core_source or 'unavailable'}") print(f"sglang_packed_provider={sglang_packed_source or 'unavailable'}") print("| tokens | native_core_ms | triton_core_ms | sglang_core_ms | fused_layout_kda_ms | sglang_packed_ms | full_ms | triton/native | sglang/native | packed/fused | native_us_per_token | triton_us_per_token | sglang_us_per_token | fused_us_per_token | packed_us_per_token | full_us_per_token |") @@ -518,9 +542,9 @@ def main() -> int: sglang_packed_ms = None full_ms = None if args.scope in ("core", "both"): - native_core_ms = bench_native_core(tokens, device, args.warmup, args.rep, args.seed) + native_core_ms = bench_native_core(tokens, device, args.warmup, args.rep, args.seed, config) if not args.skip_triton: - triton_core_ms = bench_triton_core(tokens, device, args.warmup, args.rep, args.seed) + triton_core_ms = bench_triton_core(tokens, device, args.warmup, args.rep, args.seed, config) if sglang_fused_update is not None: sglang_core_ms = bench_sglang_core( tokens, @@ -529,9 +553,10 @@ def main() -> int: args.rep, args.seed, sglang_fused_update, + config, ) if args.scope in ("fused", "both"): - fused_layout_kda_ms = bench_fused_layout_kda(tokens, device, args.warmup, args.rep, args.seed) + fused_layout_kda_ms = bench_fused_layout_kda(tokens, device, args.warmup, args.rep, args.seed, config) if sglang_packed_decode is not None: sglang_packed_ms = bench_sglang_packed_layout_kda( tokens, @@ -540,9 +565,10 @@ def main() -> int: args.rep, args.seed, sglang_packed_decode, + config, ) if args.scope in ("full", "both"): - full_ms = bench_full(tokens, device, args.warmup, args.rep, args.seed) + full_ms = bench_full(tokens, device, args.warmup, args.rep, args.seed, config) native_core_us = None if native_core_ms is None else native_core_ms * 1000.0 / tokens triton_core_us = None if triton_core_ms is None else triton_core_ms * 1000.0 / tokens @@ -580,6 +606,10 @@ def main() -> int: rows.append( { "tokens": tokens, + "tp_size": args.tp_size, + "local_k_heads": config.num_k_heads, + "local_v_heads": config.num_v_heads, + "conv_dim": config.conv_dim, "native_core_ms": native_core_ms, "triton_core_ms": triton_core_ms, "sglang_core_ms": sglang_core_ms, @@ -605,6 +635,10 @@ def main() -> int: f, fieldnames=[ "tokens", + "tp_size", + "local_k_heads", + "local_v_heads", + "conv_dim", "native_core_ms", "triton_core_ms", "sglang_core_ms", diff --git a/benchmarks/bench_qwen35_prefill.py b/benchmarks/bench_qwen35_prefill.py index ef470b7f..bfcbe6d5 100644 --- a/benchmarks/bench_qwen35_prefill.py +++ b/benchmarks/bench_qwen35_prefill.py @@ -27,8 +27,8 @@ kernel instead. Note: cula_qk currently benchmarks the TMA tensor-core Q @ K^T subpath only, -not the full gated-delta prefill recurrence. Its output is [B,48,T,T], so long -sequence lengths have quadratic memory cost. +not the full gated-delta prefill recurrence. Its output is [B,local_HV,T,T], +so long sequence lengths have quadratic memory cost. """ from __future__ import annotations @@ -51,6 +51,7 @@ from cula.ops.qwen35_fused_kda_prefill import qwen35_fused_kda_prefill from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as CONFIG +from cula.qwen35.common import Qwen35LinearAttentionConfig from cula.utils import get_kda_fused_fwd try: @@ -117,30 +118,45 @@ def resolve_sgl_chunk_gdr(sglang_path: pathlib.Path | None): return module.chunk_gated_delta_rule, "sglang.srt.layers.attention.fla.chunk.chunk_gated_delta_rule" -def make_inputs(batch: int, seq_len: int, *, device: torch.device, seed: int): +def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: + if tp_size not in (1, 2, 4, 8): + raise ValueError(f"tp_size must be one of 1, 2, 4, 8, got {tp_size}") + return Qwen35LinearAttentionConfig( + hidden_size=CONFIG.hidden_size // tp_size, + conv_kernel_size=CONFIG.conv_kernel_size, + num_k_heads=CONFIG.num_k_heads // tp_size, + num_v_heads=CONFIG.num_v_heads // tp_size, + head_k_dim=CONFIG.head_k_dim, + head_v_dim=CONFIG.head_v_dim, + qkv_dtype=CONFIG.qkv_dtype, + state_dtype=CONFIG.state_dtype, + ) + + +def make_inputs(batch: int, seq_len: int, *, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): torch.manual_seed(seed) - q = torch.randn(batch, seq_len, CONFIG.num_v_heads, CONFIG.head_k_dim, device=device, dtype=CONFIG.qkv_dtype) + q = torch.randn(batch, seq_len, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) k = torch.randn_like(q) v = torch.randn_like(q) - a = torch.randn(batch, seq_len, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - b = torch.randn(batch, seq_len, CONFIG.num_v_heads, device=device, dtype=CONFIG.qkv_dtype) - beta = torch.sigmoid(b.float()).to(dtype=CONFIG.qkv_dtype) - A_log = -torch.rand(CONFIG.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(CONFIG.num_v_heads, device=device, dtype=torch.float32) * 0.1 + a = torch.randn(batch, seq_len, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(batch, seq_len, config.num_v_heads, device=device, dtype=config.qkv_dtype) + beta = torch.sigmoid(b.float()).to(dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 log_gate = (-torch.exp(A_log).view(1, 1, -1) * torch.nn.functional.softplus(a.float() + dt_bias.view(1, 1, -1))).to( - dtype=CONFIG.qkv_dtype + dtype=config.qkv_dtype ) initial_state = torch.randn( batch, - CONFIG.num_v_heads, - CONFIG.head_k_dim, - CONFIG.head_v_dim, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, device=device, dtype=torch.float32, ) * 0.01 - mixed_qkv_conv = torch.randn(batch * seq_len, CONFIG.conv_dim, device=device, dtype=CONFIG.qkv_dtype) - a_flat = a.reshape(batch * seq_len, CONFIG.num_v_heads).contiguous() - b_flat = b.reshape(batch * seq_len, CONFIG.num_v_heads).contiguous() + mixed_qkv_conv = torch.randn(batch * seq_len, config.conv_dim, device=device, dtype=config.qkv_dtype) + a_flat = a.reshape(batch * seq_len, config.num_v_heads).contiguous() + b_flat = b.reshape(batch * seq_len, config.num_v_heads).contiguous() return q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, mixed_qkv_conv, a_flat, b_flat @@ -195,7 +211,7 @@ def prepare_cula_fused_core_inputs(q, k, a, b, A_log, dt_bias, initial_state): return q_norm, k_norm, log_gate_cumsum, beta, initial_state_vk -def run_cula_fused_core(q_norm, k_norm, v, log_gate_cumsum, beta, initial_state_vk): +def run_cula_fused_core(q_norm, k_norm, v, log_gate_cumsum, beta, initial_state_vk, config: Qwen35LinearAttentionConfig): fused_kda_prefill = get_kda_fused_fwd(q_norm.device) return fused_kda_prefill( q=q_norm, @@ -203,7 +219,7 @@ def run_cula_fused_core(q_norm, k_norm, v, log_gate_cumsum, beta, initial_state_ v=v.contiguous(), g=log_gate_cumsum, beta=beta, - scale=CONFIG.head_k_dim**-0.5, + scale=config.head_k_dim**-0.5, initial_state=initial_state_vk, output_final_state=True, use_qk_l2norm_in_kernel=False, @@ -213,7 +229,7 @@ def run_cula_fused_core(q_norm, k_norm, v, log_gate_cumsum, beta, initial_state_ ) -def run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices): +def run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices, config: Qwen35LinearAttentionConfig): # SGLang/FLA GDR chunk kernels use [N, H, V, K] state layout. cuLA's # Qwen3.5 wrapper uses [N, H, K, V], so pass the transposed view here. initial_state_vk = initial_state.transpose(-1, -2).contiguous() @@ -226,7 +242,7 @@ def run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_sta initial_state=initial_state_vk, initial_state_indices=initial_state_indices, output_final_state=True, - scale=CONFIG.head_k_dim**-0.5, + scale=config.head_k_dim**-0.5, use_qk_l2norm_in_kernel=True, head_first=False, ) @@ -253,10 +269,15 @@ def _state_to_cula_layout(state: torch.Tensor | None) -> torch.Tensor | None: def print_header(device: torch.device, args: argparse.Namespace, baseline_sources: dict[str, str]) -> None: + config = local_config_from_tp_size(args.tp_size) print("Qwen3.5 prefill benchmark") print(f" device: {torch.cuda.get_device_name(device)}") - print(f" dtype: {CONFIG.qkv_dtype}") + print(f" dtype: {config.qkv_dtype}") print(f" batch: {args.batch}") + print( + f" tp/local config: tp={args.tp_size} local_k_heads={config.num_k_heads} " + f"local_v_heads={config.num_v_heads} conv_dim={config.conv_dim}" + ) print(f" seq lens: {args.seq_lens}") print(f" warmup/rep: {args.warmup}/{args.rep}") print(f" baselines: {baseline_sources or 'disabled/unavailable'}") @@ -284,6 +305,7 @@ def main() -> None: parser.add_argument("--warmup", type=int, default=10) parser.add_argument("--rep", type=int, default=30) parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--tp-size", type=int, choices=[1, 2, 4, 8], default=1) parser.add_argument("--baseline", choices=["none", "fla", "sgl", "all"], default="sgl") parser.add_argument("--sglang-path", type=pathlib.Path, default=None) parser.add_argument( @@ -297,9 +319,10 @@ def main() -> None: "--max-qk-elements", type=int, default=512 * 1024 * 1024, - help="Skip cuLA QK timings when B*48*T*T exceeds this element count.", + help="Skip cuLA QK timings when B*local_HV*T*T exceeds this element count.", ) args = parser.parse_args() + config = local_config_from_tp_size(args.tp_size) if not torch.cuda.is_available(): raise RuntimeError("CUDA is required for this benchmark.") @@ -330,18 +353,19 @@ def main() -> None: seq_len, device=device, seed=args.seed, + config=config, ) initial_state_indices = torch.arange(args.batch, device=device, dtype=torch.int32) def layout_fn(): return qwen35_layout_prefill(mixed_qkv_conv, a_flat, b_flat, backend="cudac") - qk_elements = args.batch * CONFIG.num_v_heads * seq_len * seq_len + qk_elements = args.batch * config.num_v_heads * seq_len * seq_len qk_out = None if args.cula_mode == "qk" and qk_elements <= args.max_qk_elements: qk_out = torch.empty( args.batch, - CONFIG.num_v_heads, + config.num_v_heads, seq_len, seq_len, device=device, @@ -357,7 +381,7 @@ def cula_fn(): if args.cula_mode == "fused": return run_cula_fused(q, k, v, a, b, A_log, dt_bias, initial_state) if args.cula_mode == "fused-core": - return run_cula_fused_core(*fused_core_inputs[:2], v, *fused_core_inputs[2:]) + return run_cula_fused_core(*fused_core_inputs[:2], v, *fused_core_inputs[2:], config) if qk_out is None: raise RuntimeError( f"Skipping cuLA QK: B*H*T*T={qk_elements} exceeds --max-qk-elements={args.max_qk_elements}" @@ -394,7 +418,7 @@ def cula_fn(): for baseline_name, chunk_gdr in baselines.items(): def baseline_fn(): - return run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices) + return run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices, config) row_rel_rms = rel_rms row_rel_max = rel_max diff --git a/csrc/qwen35/decode/qwen35_conv1d_decode.cu b/csrc/qwen35/decode/qwen35_conv1d_decode.cu index 34020375..5a3c3470 100644 --- a/csrc/qwen35/decode/qwen35_conv1d_decode.cu +++ b/csrc/qwen35/decode/qwen35_conv1d_decode.cu @@ -67,20 +67,21 @@ __global__ void qwen35_conv1d_decode_kernel( scalar_t* __restrict__ conv_state, const scalar_t* __restrict__ conv_weight, scalar_t* __restrict__ out, - int batch_size) { + int batch_size, + int conv_dim) { constexpr int kThreads = 256; const int64_t linear_idx = static_cast(blockIdx.x) * kThreads + threadIdx.x; - const int64_t total = static_cast(batch_size) * cula::qwen35::decode::kMixedQKVDim; + const int64_t total = static_cast(batch_size) * conv_dim; if (linear_idx >= total) { return; } - const int64_t b = linear_idx / cula::qwen35::decode::kMixedQKVDim; - const int64_t c = linear_idx % cula::qwen35::decode::kMixedQKVDim; + const int64_t b = linear_idx / conv_dim; + const int64_t c = linear_idx % conv_dim; - const int64_t x_idx = b * cula::qwen35::decode::kMixedQKVDim + c; + const int64_t x_idx = b * conv_dim + c; const int64_t state_base = - (b * cula::qwen35::decode::kMixedQKVDim + c) * cula::qwen35::decode::kConvKernelSize; + (b * conv_dim + c) * cula::qwen35::decode::kConvKernelSize; const int64_t weight_base = c * cula::qwen35::decode::kConvKernelSize; const float s0 = to_float(conv_state[state_base + 1]); @@ -140,32 +141,34 @@ void run_qwen35_conv1d_decode(ConvDecodeParams& params) { "conv decode only supports half/bfloat16."); const int64_t batch_size = mixed_qkv.size(0); + const int64_t conv_dim = mixed_qkv.size(2); + TORCH_CHECK(conv_dim > 0, "conv_dim must be positive."); TORCH_CHECK( - mixed_qkv.dim() == 3 && mixed_qkv.sizes() == at::IntArrayRef({batch_size, 1, kMixedQKVDim}), - "mixed_qkv must have shape [B, 1, 10240]."); + mixed_qkv.dim() == 3 && mixed_qkv.sizes() == at::IntArrayRef({batch_size, 1, conv_dim}), + "mixed_qkv must have shape [B, 1, local_conv_dim]."); TORCH_CHECK( conv_state.dim() == 3 && - conv_state.sizes() == at::IntArrayRef({batch_size, kMixedQKVDim, kConvKernelSize}), - "conv_state must have shape [B, 10240, 4]."); + conv_state.sizes() == at::IntArrayRef({batch_size, conv_dim, kConvKernelSize}), + "conv_state must have shape [B, local_conv_dim, 4]."); TORCH_CHECK( - (conv_weight.dim() == 2 && conv_weight.sizes() == at::IntArrayRef({kMixedQKVDim, kConvKernelSize})) || + (conv_weight.dim() == 2 && conv_weight.sizes() == at::IntArrayRef({conv_dim, kConvKernelSize})) || (conv_weight.dim() == 3 && - conv_weight.sizes() == at::IntArrayRef({kMixedQKVDim, 1, kConvKernelSize})), - "conv_weight must have shape [10240, 4] or [10240, 1, 4]."); + conv_weight.sizes() == at::IntArrayRef({conv_dim, 1, kConvKernelSize})), + "conv_weight must have shape [local_conv_dim, 4] or [local_conv_dim, 1, 4]."); TORCH_CHECK( - out.dim() == 3 && out.sizes() == at::IntArrayRef({batch_size, 1, kMixedQKVDim}), - "out must have shape [B, 1, 10240]."); + out.dim() == 3 && out.sizes() == at::IntArrayRef({batch_size, 1, conv_dim}), + "out must have shape [B, 1, local_conv_dim]."); const at::cuda::OptionalCUDAGuard device_guard(device); cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); - const at::Tensor mixed_qkv_2d = mixed_qkv.view({batch_size, kMixedQKVDim}); - const at::Tensor out_2d = out.view({batch_size, kMixedQKVDim}); + const at::Tensor mixed_qkv_2d = mixed_qkv.view({batch_size, conv_dim}); + const at::Tensor out_2d = out.view({batch_size, conv_dim}); const at::Tensor weight_2d = - conv_weight.dim() == 3 ? conv_weight.view({kMixedQKVDim, kConvKernelSize}) : conv_weight; + conv_weight.dim() == 3 ? conv_weight.view({conv_dim, kConvKernelSize}) : conv_weight; constexpr int kThreads = 256; - const int64_t total = batch_size * static_cast(kMixedQKVDim); + const int64_t total = batch_size * conv_dim; const dim3 block(kThreads, 1, 1); const dim3 grid(static_cast((total + kThreads - 1) / kThreads), 1, 1); @@ -180,7 +183,8 @@ void run_qwen35_conv1d_decode(ConvDecodeParams& params) { conv_state.data_ptr(), weight_2d.data_ptr(), out_2d.data_ptr(), - static_cast(batch_size)); + static_cast(batch_size), + static_cast(conv_dim)); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/csrc/qwen35/decode/qwen35_decode_common.cuh b/csrc/qwen35/decode/qwen35_decode_common.cuh index 8fad1378..061691cb 100644 --- a/csrc/qwen35/decode/qwen35_decode_common.cuh +++ b/csrc/qwen35/decode/qwen35_decode_common.cuh @@ -30,22 +30,42 @@ inline constexpr int kKDim = kNumQKHeads * kHeadDimQK; inline constexpr int kVDim = kNumVHeads * kHeadDimV; inline constexpr int kMixedQKVDim = kQDim + kKDim + kVDim; +inline constexpr int local_qk_heads_from_v_heads(int local_v_heads) { + return local_v_heads / (kNumVHeads / kNumQKHeads); +} + +inline constexpr int local_q_dim(int local_qk_heads) { + return local_qk_heads * kHeadDimQK; +} + +inline constexpr int local_v_dim(int local_v_heads) { + return local_v_heads * kHeadDimV; +} + +inline constexpr int local_mixed_qkv_dim(int local_qk_heads, int local_v_heads) { + return 2 * local_q_dim(local_qk_heads) + local_v_dim(local_v_heads); +} + +inline constexpr bool is_supported_local_v_heads(int local_v_heads) { + return local_v_heads == 48 || local_v_heads == 24 || local_v_heads == 12 || local_v_heads == 6; +} + struct ConvDecodeParams { - at::Tensor mixed_qkv; // [B, 1, 10240] - at::Tensor conv_state; // [B, 10240, 4] - at::Tensor conv_weight; // [10240, 4] - at::Tensor out; // [B, 1, 10240] + at::Tensor mixed_qkv; // [B, 1, local_conv_dim] + at::Tensor conv_state; // [B, local_conv_dim, 4] + at::Tensor conv_weight; // [local_conv_dim, 4] + at::Tensor out; // [B, 1, local_conv_dim] }; struct LayoutDecodeParams { - at::Tensor mixed_qkv_conv; // [N, 10240] - at::Tensor a; // [N, 48] - at::Tensor b; // [N, 48] - at::Tensor q_rep; // [N, 48, 128] - at::Tensor k_rep; // [N, 48, 128] - at::Tensor v; // [N, 48, 128] - at::Tensor a_kernel; // [N, 48] - at::Tensor b_kernel; // [N, 48] + at::Tensor mixed_qkv_conv; // [N, local_conv_dim] + at::Tensor a; // [N, local_v_heads] + at::Tensor b; // [N, local_v_heads] + at::Tensor q_rep; // [N, local_v_heads, 128] + at::Tensor k_rep; // [N, local_v_heads, 128] + at::Tensor v; // [N, local_v_heads, 128] + at::Tensor a_kernel; // [N, local_v_heads] + at::Tensor b_kernel; // [N, local_v_heads] }; struct ScalarKdaDecodeParams { @@ -54,27 +74,27 @@ struct ScalarKdaDecodeParams { // q_rep, k_rep, v, a_kernel, b_kernel, out // - recurrent parameters / state: float32 // A_log, dt_bias, recurrent_state - at::Tensor q_rep; // [N, 48, 128] - at::Tensor k_rep; // [N, 48, 128] - at::Tensor v; // [N, 48, 128] - at::Tensor a_kernel; // [N, 48] - at::Tensor b_kernel; // [N, 48] - at::Tensor A_log; // [48], float32 - at::Tensor dt_bias; // [48], float32 - at::Tensor recurrent_state; // [pool, 48, 128, 128], float32 + at::Tensor q_rep; // [N, local_v_heads, 128] + at::Tensor k_rep; // [N, local_v_heads, 128] + at::Tensor v; // [N, local_v_heads, 128] + at::Tensor a_kernel; // [N, local_v_heads] + at::Tensor b_kernel; // [N, local_v_heads] + at::Tensor A_log; // [local_v_heads], float32 + at::Tensor dt_bias; // [local_v_heads], float32 + at::Tensor recurrent_state; // [pool, local_v_heads, 128, 128], float32 at::Tensor pool_idx; // [N], int32 - at::Tensor out; // [N, 48, 128] + at::Tensor out; // [N, local_v_heads, 128] }; struct LayoutScalarKdaDecodeParams { - at::Tensor mixed_qkv_conv; // [N, 10240] - at::Tensor a; // [N, 48] - at::Tensor b; // [N, 48] - at::Tensor A_log; // [48], float32 - at::Tensor dt_bias; // [48], float32 - at::Tensor recurrent_state; // [pool, 48, 128, 128], float32 + at::Tensor mixed_qkv_conv; // [N, local_conv_dim] + at::Tensor a; // [N, local_v_heads] + at::Tensor b; // [N, local_v_heads] + at::Tensor A_log; // [local_v_heads], float32 + at::Tensor dt_bias; // [local_v_heads], float32 + at::Tensor recurrent_state; // [pool, local_v_heads, 128, 128], float32 at::Tensor pool_idx; // [N], int32 - at::Tensor out; // [N, 48, 128] + at::Tensor out; // [N, local_v_heads, 128] }; void run_qwen35_conv1d_decode(ConvDecodeParams& params); diff --git a/csrc/qwen35/decode/qwen35_layout_decode.cu b/csrc/qwen35/decode/qwen35_layout_decode.cu index 06d74931..138730ea 100644 --- a/csrc/qwen35/decode/qwen35_layout_decode.cu +++ b/csrc/qwen35/decode/qwen35_layout_decode.cu @@ -37,6 +37,34 @@ void check_tensor_shape_2d(const at::Tensor& tensor, const char* name) { "."); } +template +void launch_layout_decode_for_heads( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t batch_size) { + constexpr int kLocalQKHeads = cula::qwen35::decode::local_qk_heads_from_v_heads(kLocalVHeads); + constexpr int threads = 32; + dim3 grid(kLocalVHeads, static_cast(batch_size), 1); + cula::qwen35::decode::qwen35_layout_decode_kernel_cute + <<>>( + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + batch_size); +} + } // namespace namespace cula::qwen35::decode { @@ -66,11 +94,15 @@ void run_qwen35_layout_decode(LayoutDecodeParams& params) { check_tensor_shape_2d(a, "a"); check_tensor_shape_2d(b, "b"); - TORCH_CHECK( - mixed_qkv_conv.dim() == 2 && mixed_qkv_conv.size(1) == kMixedQKVDim, - "mixed_qkv_conv must have shape [N, 10240]."); - const int64_t batch_size = mixed_qkv_conv.size(0); + const int64_t local_v_heads = a.size(1); + TORCH_CHECK(is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + const int local_qk_heads = local_qk_heads_from_v_heads(static_cast(local_v_heads)); + const int local_mixed_dim = local_mixed_qkv_dim(local_qk_heads, static_cast(local_v_heads)); + TORCH_CHECK( + mixed_qkv_conv.dim() == 2 && mixed_qkv_conv.size(1) == local_mixed_dim, + "mixed_qkv_conv must have shape [N, local_conv_dim=", local_mixed_dim, "], got ", + mixed_qkv_conv.sizes(), "."); const at::Device device = mixed_qkv_conv.device(); check_tensor_device(a, "a", device); @@ -88,27 +120,25 @@ void run_qwen35_layout_decode(LayoutDecodeParams& params) { TORCH_CHECK(b_kernel.is_contiguous(), "b_kernel must be contiguous."); TORCH_CHECK( - q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({batch_size, kNumVHeads, kHeadDimQK}), - "q_rep must have shape [N, 48, 128]."); + q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({batch_size, local_v_heads, kHeadDimQK}), + "q_rep must have shape [N, local_v_heads, 128]."); TORCH_CHECK( - k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({batch_size, kNumVHeads, kHeadDimQK}), - "k_rep must have shape [N, 48, 128]."); + k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({batch_size, local_v_heads, kHeadDimQK}), + "k_rep must have shape [N, local_v_heads, 128]."); TORCH_CHECK( - v.dim() == 3 && v.sizes() == at::IntArrayRef({batch_size, kNumVHeads, kHeadDimV}), - "v must have shape [N, 48, 128]."); + v.dim() == 3 && v.sizes() == at::IntArrayRef({batch_size, local_v_heads, kHeadDimV}), + "v must have shape [N, local_v_heads, 128]."); TORCH_CHECK( - a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), - "a_kernel must have shape [N, 48]."); + a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({batch_size, local_v_heads}), + "a_kernel must have shape [N, local_v_heads]."); TORCH_CHECK( - b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), - "b_kernel must have shape [N, 48]."); + b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({batch_size, local_v_heads}), + "b_kernel must have shape [N, local_v_heads]."); - TORCH_CHECK(a.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), "a must have shape [N, 48]."); - TORCH_CHECK(b.sizes() == at::IntArrayRef({batch_size, kNumVHeads}), "b must have shape [N, 48]."); + TORCH_CHECK(a.sizes() == at::IntArrayRef({batch_size, local_v_heads}), "a must have shape [N, local_v_heads]."); + TORCH_CHECK(b.sizes() == at::IntArrayRef({batch_size, local_v_heads}), "b must have shape [N, local_v_heads]."); const at::cuda::OptionalCUDAGuard device_guard(device); - constexpr int threads = 32; - dim3 grid(kNumVHeads, static_cast(batch_size), 1); cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); AT_DISPATCH_FLOATING_TYPES_AND2( @@ -117,16 +147,20 @@ void run_qwen35_layout_decode(LayoutDecodeParams& params) { mixed_qkv_conv.scalar_type(), "qwen35_layout_decode_kernel_cute", [&] { - qwen35_layout_decode_kernel_cute<<>>( - mixed_qkv_conv.data_ptr(), - a.data_ptr(), - b.data_ptr(), - q_rep.data_ptr(), - k_rep.data_ptr(), - v.data_ptr(), - a_kernel.data_ptr(), - b_kernel.data_ptr(), - batch_size); + switch (local_v_heads) { + case 48: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + case 24: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + case 12: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + case 6: + launch_layout_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), batch_size); + break; + } }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/csrc/qwen35/decode/qwen35_layout_kernel.hpp b/csrc/qwen35/decode/qwen35_layout_kernel.hpp index 36e9373e..b665dce2 100644 --- a/csrc/qwen35/decode/qwen35_layout_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_layout_kernel.hpp @@ -45,7 +45,7 @@ CUTE_DEVICE void copy_vec_contiguous( } } -template +template __global__ void qwen35_layout_decode_kernel_cute( const scalar_t* __restrict__ mixed_qkv_conv, const scalar_t* __restrict__ a, @@ -56,8 +56,11 @@ __global__ void qwen35_layout_decode_kernel_cute( scalar_t* __restrict__ a_kernel, scalar_t* __restrict__ b_kernel, int64_t token_count) { - static_assert(kNumVHeads % kNumQKHeads == 0); - constexpr int kRepeatFactor = kNumVHeads / kNumQKHeads; + static_assert(kLocalVHeads % kLocalQKHeads == 0); + constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; + constexpr int kLocalQDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalKDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalMixedQKVDim = 2 * kLocalQDim + kLocalVHeads * kHeadDimV; // TODO(qwen35-layout-opt): // - Re-evaluate whether Vec=8 is profitable for bf16/fp16 on the target GPUs. // - Push more of the q/k repeat mapping into compile-time CuTe layout transforms. @@ -72,31 +75,31 @@ __global__ void qwen35_layout_decode_kernel_cute( const int hv = static_cast(blockIdx.x); const int tid = static_cast(threadIdx.x); - if (token_idx >= token_count || hv >= kNumVHeads) { + if (token_idx >= token_count || hv >= kLocalVHeads) { return; } const int mapped_h = hv / kRepeatFactor; auto qk_src_layout = make_layout( - make_shape(Int{}, Int{}), + make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{})); auto v_src_layout = make_layout( - make_shape(Int{}, Int{}), + make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{})); auto out_layout = make_layout( - make_shape(Int{}, Int{}), + make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{})); - auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); - const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kMixedQKVDim; + const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kLocalMixedQKVDim; const scalar_t* q_src_ptr = token_ptr; - const scalar_t* k_src_ptr = token_ptr + kQDim; - const scalar_t* v_src_ptr = token_ptr + kQDim + kKDim; + const scalar_t* k_src_ptr = token_ptr + kLocalQDim; + const scalar_t* v_src_ptr = token_ptr + kLocalQDim + kLocalKDim; - scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; - scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; - scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kNumVHeads * kHeadDimV; + scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kLocalVHeads * kHeadDimV; // Current version uses a direct GMEM->GMEM vector copy path. This keeps the // kernel simple while already removing the scalar-copy bottleneck from the @@ -118,7 +121,7 @@ __global__ void qwen35_layout_decode_kernel_cute( // TODO(qwen35-layout-opt): If a/b copy becomes measurable, fuse a wider // per-head copy path here instead of scalar head writes. const int head_idx = crd2idx(make_coord(hv), head_layout); - const int64_t token_head_offset = static_cast(token_idx) * kNumVHeads + head_idx; + const int64_t token_head_offset = static_cast(token_idx) * kLocalVHeads + head_idx; a_kernel[token_head_offset] = a[token_head_offset]; b_kernel[token_head_offset] = b[token_head_offset]; } diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu index a69f5bc6..8718b4ff 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu +++ b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu @@ -28,6 +28,62 @@ void check_tensor_device(const at::Tensor& tensor, const char* name, const at::D TORCH_CHECK(tensor.device() == device, name, " must be on device ", device, "."); } +template +void dispatch_scalar_decode_for_heads( + cudaStream_t stream, + const scalar_t* q_rep, + const scalar_t* k_rep, + const scalar_t* v, + const scalar_t* a_kernel, + const scalar_t* b_kernel, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + constexpr int kLocalQKHeads = local_qk_heads_from_v_heads(kLocalVHeads); + kernel::launch_qwen35_scalar_kda_decode_kernel( + stream, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + +template +void dispatch_layout_scalar_decode_for_heads( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + constexpr int kLocalQKHeads = local_qk_heads_from_v_heads(kLocalVHeads); + kernel::launch_qwen35_layout_scalar_kda_decode_kernel( + stream, + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + } // namespace void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params) { @@ -76,34 +132,37 @@ void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params) { TORCH_CHECK(recurrent_state.scalar_type() == at::kFloat, "recurrent_state must be float32."); TORCH_CHECK(pool_idx.scalar_type() == at::kInt, "pool_idx must be int32."); + TORCH_CHECK(q_rep.dim() == 3, "q_rep must have shape [N, local_v_heads, 128]."); const int64_t token_count = q_rep.size(0); + const int64_t local_v_heads = q_rep.size(1); + TORCH_CHECK(is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); TORCH_CHECK( - q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimQK}), - "q_rep must have shape [N, 48, 128]."); + q_rep.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimQK}), + "q_rep must have shape [N, local_v_heads, 128]."); TORCH_CHECK( - k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimQK}), - "k_rep must have shape [N, 48, 128]."); + k_rep.dim() == 3 && k_rep.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimQK}), + "k_rep must have shape [N, local_v_heads, 128]."); TORCH_CHECK( - v.dim() == 3 && v.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimV}), - "v must have shape [N, 48, 128]."); + v.dim() == 3 && v.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimV}), + "v must have shape [N, local_v_heads, 128]."); TORCH_CHECK( - a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({token_count, kNumVHeads}), - "a_kernel must have shape [N, 48]."); + a_kernel.dim() == 2 && a_kernel.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "a_kernel must have shape [N, local_v_heads]."); TORCH_CHECK( - b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({token_count, kNumVHeads}), - "b_kernel must have shape [N, 48]."); - TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == kNumVHeads, "A_log must have shape [48]."); - TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == kNumVHeads, "dt_bias must have shape [48]."); + b_kernel.dim() == 2 && b_kernel.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "b_kernel must have shape [N, local_v_heads]."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == local_v_heads, "A_log must have shape [local_v_heads]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == local_v_heads, "dt_bias must have shape [local_v_heads]."); TORCH_CHECK( recurrent_state.dim() == 4 && - recurrent_state.size(1) == kNumVHeads && + recurrent_state.size(1) == local_v_heads && recurrent_state.size(2) == kHeadDimQK && recurrent_state.size(3) == kHeadDimV, - "recurrent_state must have shape [pool, 48, 128, 128]."); + "recurrent_state must have shape [pool, local_v_heads, 128, 128]."); TORCH_CHECK(pool_idx.dim() == 1 && pool_idx.size(0) == token_count, "pool_idx must have shape [N]."); TORCH_CHECK( - out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimV}), - "out must have shape [N, 48, 128]."); + out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimV}), + "out must have shape [N, local_v_heads, 128]."); const at::cuda::OptionalCUDAGuard device_guard(device); cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); @@ -114,19 +173,20 @@ void run_qwen35_scalar_kda_decode(ScalarKdaDecodeParams& params) { q_rep.scalar_type(), "launch_qwen35_scalar_kda_decode_kernel", [&] { - kernel::launch_qwen35_scalar_kda_decode_kernel( - stream, - q_rep.data_ptr(), - k_rep.data_ptr(), - v.data_ptr(), - a_kernel.data_ptr(), - b_kernel.data_ptr(), - A_log.data_ptr(), - dt_bias.data_ptr(), - recurrent_state.data_ptr(), - pool_idx.data_ptr(), - out.data_ptr(), - static_cast(token_count)); + switch (local_v_heads) { + case 48: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 24: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 12: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 6: + dispatch_scalar_decode_for_heads(stream, q_rep.data_ptr(), k_rep.data_ptr(), v.data_ptr(), a_kernel.data_ptr(), b_kernel.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + } }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -174,29 +234,34 @@ void run_qwen35_layout_scalar_kda_decode(LayoutScalarKdaDecodeParams& params) { TORCH_CHECK(recurrent_state.scalar_type() == at::kFloat, "recurrent_state must be float32."); TORCH_CHECK(pool_idx.scalar_type() == at::kInt, "pool_idx must be int32."); + TORCH_CHECK(mixed_qkv_conv.dim() == 2, "mixed_qkv_conv must have shape [N, local_conv_dim]."); + TORCH_CHECK(a.dim() == 2, "a must have shape [N, local_v_heads]."); const int64_t token_count = mixed_qkv_conv.size(0); + const int64_t local_v_heads = a.size(1); + TORCH_CHECK(is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + const int local_qk_heads = local_qk_heads_from_v_heads(static_cast(local_v_heads)); + const int local_mixed_dim = local_mixed_qkv_dim(local_qk_heads, static_cast(local_v_heads)); TORCH_CHECK( - mixed_qkv_conv.dim() == 2 && - mixed_qkv_conv.sizes() == at::IntArrayRef({token_count, kMixedQKVDim}), - "mixed_qkv_conv must have shape [N, 10240]."); + mixed_qkv_conv.sizes() == at::IntArrayRef({token_count, local_mixed_dim}), + "mixed_qkv_conv must have shape [N, local_conv_dim]."); TORCH_CHECK( - a.dim() == 2 && a.sizes() == at::IntArrayRef({token_count, kNumVHeads}), - "a must have shape [N, 48]."); + a.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "a must have shape [N, local_v_heads]."); TORCH_CHECK( - b.dim() == 2 && b.sizes() == at::IntArrayRef({token_count, kNumVHeads}), - "b must have shape [N, 48]."); - TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == kNumVHeads, "A_log must have shape [48]."); - TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == kNumVHeads, "dt_bias must have shape [48]."); + b.dim() == 2 && b.sizes() == at::IntArrayRef({token_count, local_v_heads}), + "b must have shape [N, local_v_heads]."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == local_v_heads, "A_log must have shape [local_v_heads]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == local_v_heads, "dt_bias must have shape [local_v_heads]."); TORCH_CHECK( recurrent_state.dim() == 4 && - recurrent_state.size(1) == kNumVHeads && + recurrent_state.size(1) == local_v_heads && recurrent_state.size(2) == kHeadDimQK && recurrent_state.size(3) == kHeadDimV, - "recurrent_state must have shape [pool, 48, 128, 128]."); + "recurrent_state must have shape [pool, local_v_heads, 128, 128]."); TORCH_CHECK(pool_idx.dim() == 1 && pool_idx.size(0) == token_count, "pool_idx must have shape [N]."); TORCH_CHECK( - out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimV}), - "out must have shape [N, 48, 128]."); + out.dim() == 3 && out.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimV}), + "out must have shape [N, local_v_heads, 128]."); const at::cuda::OptionalCUDAGuard device_guard(device); cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); @@ -207,17 +272,20 @@ void run_qwen35_layout_scalar_kda_decode(LayoutScalarKdaDecodeParams& params) { mixed_qkv_conv.scalar_type(), "launch_qwen35_layout_scalar_kda_decode_kernel", [&] { - kernel::launch_qwen35_layout_scalar_kda_decode_kernel( - stream, - mixed_qkv_conv.data_ptr(), - a.data_ptr(), - b.data_ptr(), - A_log.data_ptr(), - dt_bias.data_ptr(), - recurrent_state.data_ptr(), - pool_idx.data_ptr(), - out.data_ptr(), - static_cast(token_count)); + switch (local_v_heads) { + case 48: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 24: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 12: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + case 6: + dispatch_layout_scalar_decode_for_heads(stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), b.data_ptr(), A_log.data_ptr(), dt_bias.data_ptr(), recurrent_state.data_ptr(), pool_idx.data_ptr(), out.data_ptr(), static_cast(token_count)); + break; + } }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 1a767524..3cd032f4 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -23,7 +23,7 @@ namespace cula::qwen35::decode::kernel { using namespace cute; -template +template struct Qwen35ScalarKdaDecodeKernel { // Decode-first design: // - 1 CTA owns 1 (token_idx, hv) @@ -39,7 +39,7 @@ struct Qwen35ScalarKdaDecodeKernel { static constexpr int kTilesPerV = kHeadDimV / kTileV; static constexpr int kTilesPerK = kHeadDimQK / kTileK; - static_assert(kNumQKHeads < kNumVHeads); + static_assert(kLocalQKHeads < kLocalVHeads); static_assert(kHeadDimQK == 128); static_assert(kHeadDimV == 128); static_assert(kThreads == kWarpGroupThreads); @@ -65,7 +65,7 @@ struct Qwen35ScalarKdaDecodeKernel { static dim3 grid_shape(int token_count) { // One block owns one (token_idx, hv) pair in the first implementation. - return dim3(static_cast(kNumVHeads), static_cast(token_count), 1); + return dim3(static_cast(kLocalVHeads), static_cast(token_count), 1); } template @@ -85,7 +85,7 @@ struct Qwen35ScalarKdaDecodeKernel { const int hv = static_cast(blockIdx.x); const int token_idx = static_cast(blockIdx.y); const int tid = static_cast(threadIdx.x); - if (token_idx >= token_count || hv >= kNumVHeads) { + if (token_idx >= token_count || hv >= kLocalVHeads) { return; } @@ -127,21 +127,21 @@ struct Qwen35ScalarKdaDecodeKernel { // of introducing that complexity before the math path itself is stable. auto q_layout = make_layout( - make_shape(token_count, Int{}, Int{}), - make_stride(kNumVHeads * kHeadDimQK, kHeadDimQK, Int<1>{})); + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimQK, kHeadDimQK, Int<1>{})); auto v_layout = make_layout( - make_shape(token_count, Int{}, Int{}), - make_stride(kNumVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); auto head_layout = make_layout( - make_shape(token_count, Int{}), - make_stride(kNumVHeads, Int<1>{})); - auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); auto state_layout_kv = make_layout( - make_shape(_, Int{}, Int{}, Int{}), - make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); auto state_layout_vk = make_layout( - make_shape(_, Int{}, Int{}, Int{}), - make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); auto gQ = make_tensor(make_gmem_ptr(q_rep), q_layout); auto gK = make_tensor(make_gmem_ptr(k_rep), q_layout); @@ -197,40 +197,43 @@ struct Qwen35ScalarKdaDecodeKernel { scalar_t* __restrict__ out, int token_count, SharedStorage& storage) { - constexpr int kRepeatFactor = kNumVHeads / kNumQKHeads; + constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; + constexpr int kLocalQDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalKDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalMixedQKVDim = 2 * kLocalQDim + kLocalVHeads * kHeadDimV; const int hv = static_cast(blockIdx.x); const int token_idx = static_cast(blockIdx.y); const int tid = static_cast(threadIdx.x); - if (token_idx >= token_count || hv >= kNumVHeads) { + if (token_idx >= token_count || hv >= kLocalVHeads) { return; } const int mapped_h = hv / kRepeatFactor; auto qk_src_layout = make_layout( - make_shape(token_count, Int{}, Int{}), - make_stride(kMixedQKVDim, kHeadDimQK, Int<1>{})); + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimQK, Int<1>{})); auto v_src_layout = make_layout( - make_shape(token_count, Int{}, Int{}), - make_stride(kMixedQKVDim, kHeadDimV, Int<1>{})); + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimV, Int<1>{})); auto out_layout = make_layout( - make_shape(token_count, Int{}, Int{}), - make_stride(kNumVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); auto head_layout = make_layout( - make_shape(token_count, Int{}), - make_stride(kNumVHeads, Int<1>{})); - auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); auto state_layout_kv = make_layout( - make_shape(_, Int{}, Int{}, Int{}), - make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, kHeadDimV, Int<1>{})); auto state_layout_vk = make_layout( - make_shape(_, Int{}, Int{}, Int{}), - make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); const scalar_t* q_src = mixed_qkv_conv; - const scalar_t* k_src = mixed_qkv_conv + kQDim; - const scalar_t* v_src = mixed_qkv_conv + kQDim + kKDim; + const scalar_t* k_src = mixed_qkv_conv + kLocalQDim; + const scalar_t* v_src = mixed_qkv_conv + kLocalQDim + kLocalKDim; auto gQ = make_tensor(make_gmem_ptr(q_src), qk_src_layout); auto gK = make_tensor(make_gmem_ptr(k_src), qk_src_layout); @@ -275,7 +278,7 @@ struct Qwen35ScalarKdaDecodeKernel { } }; -template > +template > __global__ void qwen35_scalar_kda_decode_kernel( const scalar_t* __restrict__ q_rep, const scalar_t* __restrict__ k_rep, @@ -288,8 +291,8 @@ __global__ void qwen35_scalar_kda_decode_kernel( const int32_t* __restrict__ pool_idx, scalar_t* __restrict__ out, int token_count) { - __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; - Qwen35ScalarKdaDecodeKernel::template run_device( + __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; + Qwen35ScalarKdaDecodeKernel::template run_device( q_rep, k_rep, v, @@ -304,7 +307,7 @@ __global__ void qwen35_scalar_kda_decode_kernel( storage); } -template > +template > void launch_qwen35_scalar_kda_decode_kernel( cudaStream_t stream, const scalar_t* q_rep, @@ -318,9 +321,9 @@ void launch_qwen35_scalar_kda_decode_kernel( const int32_t* pool_idx, scalar_t* out, int token_count) { - auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); - auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); - qwen35_scalar_kda_decode_kernel<<>>( + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_scalar_kda_decode_kernel<<>>( q_rep, k_rep, v, @@ -334,7 +337,7 @@ void launch_qwen35_scalar_kda_decode_kernel( token_count); } -template > +template > __global__ void qwen35_layout_scalar_kda_decode_kernel( const scalar_t* __restrict__ mixed_qkv_conv, const scalar_t* __restrict__ a, @@ -345,8 +348,8 @@ __global__ void qwen35_layout_scalar_kda_decode_kernel( const int32_t* __restrict__ pool_idx, scalar_t* __restrict__ out, int token_count) { - __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; - Qwen35ScalarKdaDecodeKernel::template run_layout_device( + __shared__ typename Qwen35ScalarKdaDecodeKernel::SharedStorage storage; + Qwen35ScalarKdaDecodeKernel::template run_layout_device( mixed_qkv_conv, a, b, @@ -359,7 +362,7 @@ __global__ void qwen35_layout_scalar_kda_decode_kernel( storage); } -template > +template > void launch_qwen35_layout_scalar_kda_decode_kernel( cudaStream_t stream, const scalar_t* mixed_qkv_conv, @@ -371,9 +374,9 @@ void launch_qwen35_layout_scalar_kda_decode_kernel( const int32_t* pool_idx, scalar_t* out, int token_count) { - auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); - auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); - qwen35_layout_scalar_kda_decode_kernel<<>>( + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_layout_scalar_kda_decode_kernel<<>>( mixed_qkv_conv, a, b, diff --git a/csrc/qwen35/prefill/qwen35_layout_prefill.cu b/csrc/qwen35/prefill/qwen35_layout_prefill.cu index dbda9577..5cdf1218 100644 --- a/csrc/qwen35/prefill/qwen35_layout_prefill.cu +++ b/csrc/qwen35/prefill/qwen35_layout_prefill.cu @@ -32,6 +32,61 @@ void check_rank_2(const at::Tensor& tensor, const char* name) { TORCH_CHECK(tensor.dim() == 2, name, " must be rank 2, got rank ", tensor.dim(), "."); } +template +void dispatch_layout_prefill_for_heads( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t token_count) { + constexpr int kLocalQKHeads = decode::local_qk_heads_from_v_heads(kLocalVHeads); + kernel::launch_qwen35_layout_prefill_kernel( + stream, + mixed_qkv_conv, + a, + b, + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + token_count); +} + +template +void dispatch_layout_prefill( + int64_t local_v_heads, + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + scalar_t* q_rep, + scalar_t* k_rep, + scalar_t* v, + scalar_t* a_kernel, + scalar_t* b_kernel, + int64_t token_count) { + switch (local_v_heads) { + case 48: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + case 24: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + case 12: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + case 6: + dispatch_layout_prefill_for_heads(stream, mixed_qkv_conv, a, b, q_rep, k_rep, v, a_kernel, b_kernel, token_count); + break; + } +} + } // namespace void run_qwen35_layout_prefill(LayoutPrefillParams& params) { @@ -82,12 +137,16 @@ void run_qwen35_layout_prefill(LayoutPrefillParams& params) { check_rank_2(b, "b"); const int64_t token_count = mixed_qkv_conv.size(0); - TORCH_CHECK(mixed_qkv_conv.size(1) == kMixedQKVDim, "mixed_qkv_conv must be [N, 10240]."); - TORCH_CHECK(a.sizes() == at::IntArrayRef({token_count, kNumVHeads}), "a must be [N, 48]."); - TORCH_CHECK(b.sizes() == at::IntArrayRef({token_count, kNumVHeads}), "b must be [N, 48]."); + const int64_t local_v_heads = a.size(1); + TORCH_CHECK(decode::is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + const int local_qk_heads = decode::local_qk_heads_from_v_heads(static_cast(local_v_heads)); + const int local_mixed_dim = decode::local_mixed_qkv_dim(local_qk_heads, static_cast(local_v_heads)); + TORCH_CHECK(mixed_qkv_conv.size(1) == local_mixed_dim, "mixed_qkv_conv must be [N, local_conv_dim=", local_mixed_dim, "]."); + TORCH_CHECK(a.sizes() == at::IntArrayRef({token_count, local_v_heads}), "a must be [N, local_v_heads]."); + TORCH_CHECK(b.sizes() == at::IntArrayRef({token_count, local_v_heads}), "b must be [N, local_v_heads]."); TORCH_CHECK( - q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({token_count, kNumVHeads, kHeadDimQK}), - "q_rep must be [N, 48, 128]."); + q_rep.dim() == 3 && q_rep.sizes() == at::IntArrayRef({token_count, local_v_heads, kHeadDimQK}), + "q_rep must be [N, local_v_heads, 128]."); TORCH_CHECK(k_rep.sizes() == q_rep.sizes(), "k_rep must match q_rep shape."); TORCH_CHECK(v.sizes() == q_rep.sizes(), "v must match q_rep shape."); TORCH_CHECK(a_kernel.sizes() == a.sizes(), "a_kernel must match a shape."); @@ -97,7 +156,8 @@ void run_qwen35_layout_prefill(LayoutPrefillParams& params) { cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); if (mixed_qkv_conv.scalar_type() == at::kHalf) { - kernel::launch_qwen35_layout_prefill_kernel( + dispatch_layout_prefill( + local_v_heads, stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), @@ -109,7 +169,8 @@ void run_qwen35_layout_prefill(LayoutPrefillParams& params) { b_kernel.data_ptr(), token_count); } else { - kernel::launch_qwen35_layout_prefill_kernel( + dispatch_layout_prefill( + local_v_heads, stream, mixed_qkv_conv.data_ptr(), a.data_ptr(), diff --git a/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp index b2e0adaf..7dbcae3d 100644 --- a/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp +++ b/csrc/qwen35/prefill/qwen35_layout_prefill_kernel.hpp @@ -46,7 +46,7 @@ CUTE_DEVICE void copy_prefill_vec_contiguous( } } -template +template __global__ void qwen35_layout_prefill_kernel( const scalar_t* __restrict__ mixed_qkv_conv, const scalar_t* __restrict__ a, @@ -57,40 +57,43 @@ __global__ void qwen35_layout_prefill_kernel( scalar_t* __restrict__ a_kernel, scalar_t* __restrict__ b_kernel, int64_t token_count) { - static_assert(kNumVHeads % kNumQKHeads == 0); + static_assert(kLocalVHeads % kLocalQKHeads == 0); static_assert(kHeadDimQK == kHeadDimV); - constexpr int kRepeatFactor = kNumVHeads / kNumQKHeads; + constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; + constexpr int kLocalQDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalKDim = kLocalQKHeads * kHeadDimQK; + constexpr int kLocalMixedQKVDim = 2 * kLocalQDim + kLocalVHeads * kHeadDimV; constexpr int kVec = 4; static_assert(kHeadDimQK % kVec == 0); const int hv = static_cast(blockIdx.x); const int token_idx = static_cast(blockIdx.y); const int tid = static_cast(threadIdx.x); - if (token_idx >= token_count || hv >= kNumVHeads) { + if (token_idx >= token_count || hv >= kLocalVHeads) { return; } const int mapped_h = hv / kRepeatFactor; auto qk_src_layout = make_layout( - make_shape(Int{}, Int{}), + make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{})); auto v_src_layout = make_layout( - make_shape(Int{}, Int{}), + make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{})); auto hv_layout = make_layout( - make_shape(Int{}, Int{}), + make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{})); - auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto head_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); - const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kMixedQKVDim; + const scalar_t* token_ptr = mixed_qkv_conv + static_cast(token_idx) * kLocalMixedQKVDim; const scalar_t* q_src_ptr = token_ptr; - const scalar_t* k_src_ptr = token_ptr + kQDim; - const scalar_t* v_src_ptr = token_ptr + kQDim + kKDim; + const scalar_t* k_src_ptr = token_ptr + kLocalQDim; + const scalar_t* v_src_ptr = token_ptr + kLocalQDim + kLocalKDim; - scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; - scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kNumVHeads * kHeadDimQK; - scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kNumVHeads * kHeadDimV; + scalar_t* q_dst_ptr = q_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* k_dst_ptr = k_rep + static_cast(token_idx) * kLocalVHeads * kHeadDimQK; + scalar_t* v_dst_ptr = v_out + static_cast(token_idx) * kLocalVHeads * kHeadDimV; for (int vec_idx = tid; vec_idx < kHeadDimQK / kVec; vec_idx += blockDim.x) { const int d = vec_idx * kVec; @@ -106,13 +109,13 @@ __global__ void qwen35_layout_prefill_kernel( if (tid == 0) { const int head_idx = crd2idx(make_coord(hv), head_layout); - const int64_t token_head_offset = static_cast(token_idx) * kNumVHeads + head_idx; + const int64_t token_head_offset = static_cast(token_idx) * kLocalVHeads + head_idx; a_kernel[token_head_offset] = a[token_head_offset]; b_kernel[token_head_offset] = b[token_head_offset]; } } -template +template void launch_qwen35_layout_prefill_kernel( cudaStream_t stream, const scalar_t* mixed_qkv_conv, @@ -125,8 +128,8 @@ void launch_qwen35_layout_prefill_kernel( scalar_t* b_kernel, int64_t token_count) { constexpr int kThreads = 32; - dim3 grid(kNumVHeads, static_cast(token_count), 1); - qwen35_layout_prefill_kernel<<>>( + dim3 grid(kLocalVHeads, static_cast(token_count), 1); + qwen35_layout_prefill_kernel<<>>( mixed_qkv_conv, a, b, diff --git a/csrc/qwen35/prefill/qwen35_prefill_common.cuh b/csrc/qwen35/prefill/qwen35_prefill_common.cuh index 6c7c6bf7..ba3f3b70 100644 --- a/csrc/qwen35/prefill/qwen35_prefill_common.cuh +++ b/csrc/qwen35/prefill/qwen35_prefill_common.cuh @@ -30,28 +30,28 @@ using decode::kQDim; using decode::kVDim; struct LayoutPrefillParams { - at::Tensor mixed_qkv_conv; // [N, 10240] - at::Tensor a; // [N, 48] - at::Tensor b; // [N, 48] - at::Tensor q_rep; // [N, 48, 128] - at::Tensor k_rep; // [N, 48, 128] - at::Tensor v; // [N, 48, 128] - at::Tensor a_kernel; // [N, 48] - at::Tensor b_kernel; // [N, 48] + at::Tensor mixed_qkv_conv; // [N, local_conv_dim] + at::Tensor a; // [N, local_v_heads] + at::Tensor b; // [N, local_v_heads] + at::Tensor q_rep; // [N, local_v_heads, 128] + at::Tensor k_rep; // [N, local_v_heads, 128] + at::Tensor v; // [N, local_v_heads, 128] + at::Tensor a_kernel; // [N, local_v_heads] + at::Tensor b_kernel; // [N, local_v_heads] }; struct ScalarKdaPrefillParams { - at::Tensor q; // [B, T, 48, 128] - at::Tensor k; // [B, T, 48, 128] - at::Tensor v; // [B, T, 48, 128] - at::Tensor a; // [B, T, 48] - at::Tensor b; // [B, T, 48] - at::Tensor A_log; // [48], float32 - at::Tensor dt_bias; // [48], float32 - at::Tensor initial_state; // [N, 48, 128, 128], float32, may be empty + at::Tensor q; // [B, T, local_v_heads, 128] + at::Tensor k; // [B, T, local_v_heads, 128] + at::Tensor v; // [B, T, local_v_heads, 128] + at::Tensor a; // [B, T, local_v_heads] + at::Tensor b; // [B, T, local_v_heads] + at::Tensor A_log; // [local_v_heads], float32 + at::Tensor dt_bias; // [local_v_heads], float32 + at::Tensor initial_state; // [N, local_v_heads, 128, 128], float32, may be empty at::Tensor cu_seqlens; // [N + 1], int32, may be empty - at::Tensor out; // [B, T, 48, 128] - at::Tensor final_state; // [N, 48, 128, 128], float32 + at::Tensor out; // [B, T, local_v_heads, 128] + at::Tensor final_state; // [N, local_v_heads, 128, 128], float32 }; void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params); diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu index 454d61ae..a92481a9 100644 --- a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu @@ -36,6 +36,81 @@ void check_contiguous(const at::Tensor& tensor, const char* name) { } } +template +void dispatch_scalar_prefill_for_heads( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + kernel::launch_qwen35_scalar_kda_prefill_kernel( + stream, + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + cu_seqlens, + out, + final_state, + batch_size, + seq_len, + sequence_count, + is_varlen, + has_initial_state); +} + +template +void dispatch_scalar_prefill( + int64_t local_v_heads, + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + switch (local_v_heads) { + case 48: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + break; + case 24: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + break; + case 12: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + break; + case 6: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + break; + } +} + } // namespace void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { @@ -96,15 +171,17 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { TORCH_CHECK(q.dim() == 4, "q must be [B, T, 48, 128]."); const int64_t B = q.size(0); const int64_t T = q.size(1); + const int64_t local_v_heads = q.size(2); + TORCH_CHECK(decode::is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); TORCH_CHECK( - q.sizes() == at::IntArrayRef({B, T, kNumVHeads, kHeadDimQK}), - "q must have shape [B, T, 48, 128]."); + q.sizes() == at::IntArrayRef({B, T, local_v_heads, kHeadDimQK}), + "q must have shape [B, T, local_v_heads, 128]."); TORCH_CHECK(k.sizes() == q.sizes(), "k must match q shape."); TORCH_CHECK(v.sizes() == q.sizes(), "v must match q shape."); - TORCH_CHECK(a.dim() == 3 && a.sizes() == at::IntArrayRef({B, T, kNumVHeads}), "a must be [B, T, 48]."); + TORCH_CHECK(a.dim() == 3 && a.sizes() == at::IntArrayRef({B, T, local_v_heads}), "a must be [B, T, local_v_heads]."); TORCH_CHECK(b.sizes() == a.sizes(), "b must match a shape."); - TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == kNumVHeads, "A_log must be [48]."); - TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == kNumVHeads, "dt_bias must be [48]."); + TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == local_v_heads, "A_log must be [local_v_heads]."); + TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == local_v_heads, "dt_bias must be [local_v_heads]."); TORCH_CHECK(out.sizes() == q.sizes(), "out must match q shape."); const bool is_varlen = cu_seqlens.defined() && cu_seqlens.numel() > 0; @@ -116,8 +193,8 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { TORCH_CHECK( final_state.dim() == 4 && - final_state.sizes() == at::IntArrayRef({sequence_count, kNumVHeads, kHeadDimQK, kHeadDimV}), - "final_state must be [N, 48, 128, 128]."); + final_state.sizes() == at::IntArrayRef({sequence_count, local_v_heads, kHeadDimQK, kHeadDimV}), + "final_state must be [N, local_v_heads, 128, 128]."); const bool has_initial_state = initial_state.defined() && initial_state.numel() > 0; if (has_initial_state) { TORCH_CHECK(initial_state.sizes() == final_state.sizes(), "initial_state must match final_state shape."); @@ -127,7 +204,8 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); if (q.scalar_type() == at::kHalf) { - kernel::launch_qwen35_scalar_kda_prefill_kernel( + dispatch_scalar_prefill( + local_v_heads, stream, q.data_ptr(), k.data_ptr(), @@ -146,7 +224,8 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { is_varlen, has_initial_state); } else { - kernel::launch_qwen35_scalar_kda_prefill_kernel( + dispatch_scalar_prefill( + local_v_heads, stream, q.data_ptr(), k.data_ptr(), diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp index db135475..c8b2e396 100644 --- a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp @@ -23,7 +23,7 @@ namespace cula::qwen35::prefill::kernel { using namespace cute; -template +template struct Qwen35ScalarKdaPrefillKernel { static constexpr int kThreads = 128; static constexpr int kHeadDim = kHeadDimQK; @@ -48,8 +48,8 @@ struct Qwen35ScalarKdaPrefillKernel { CUTE_HOST_DEVICE static auto make_v_work_tiles(int sequence_count) { auto problem_layout = make_layout( - make_shape(Int{}, Int{}, sequence_count), - make_stride(Int<1>{}, Int{}, Int{})); + make_shape(Int{}, Int{}, sequence_count), + make_stride(Int<1>{}, Int{}, Int{})); return zipped_divide(problem_layout, make_shape(Int{}, Int<1>{}, Int<1>{})); } @@ -80,7 +80,9 @@ struct Qwen35ScalarKdaPrefillKernel { } __syncthreads(); } - return storage.scratch[0]; + const float result = storage.scratch[0]; + __syncthreads(); + return result; } CUTE_DEVICE static void run_device( @@ -110,13 +112,13 @@ struct Qwen35ScalarKdaPrefillKernel { const int v_base = v_tile_idx * kVTile; const int tid = static_cast(threadIdx.x); - if (hv >= kNumVHeads || seq_idx >= sequence_count) { + if (hv >= kLocalVHeads || seq_idx >= sequence_count) { return; } const int token_begin = is_varlen ? static_cast(cu_seqlens[seq_idx]) : seq_idx * seq_len; const int token_end = is_varlen ? static_cast(cu_seqlens[seq_idx + 1]) : token_begin + seq_len; - const int state_base = ((seq_idx * kNumVHeads + hv) * kHeadDimQK) * kHeadDimV; + const int state_base = ((seq_idx * kLocalVHeads + hv) * kHeadDimQK) * kHeadDimV; const int kk = tid; float state_vals[kVTile]; @@ -138,8 +140,8 @@ struct Qwen35ScalarKdaPrefillKernel { for (int token = token_begin; token < token_end; ++token) { const int local_t = is_varlen ? token : token - token_begin; - const int qkv_base = ((token * kNumVHeads + hv) * kHeadDimQK); - const int gate_base = token * kNumVHeads + hv; + const int qkv_base = ((token * kLocalVHeads + hv) * kHeadDimQK); + const int gate_base = token * kLocalVHeads + hv; const float q_val = kk < kHeadDimQK ? load_as_float(q[qkv_base + kk]) : 0.0f; const float k_val = kk < kHeadDimQK ? load_as_float(k[qkv_base + kk]) : 0.0f; @@ -174,7 +176,7 @@ struct Qwen35ScalarKdaPrefillKernel { if (tid == 0) { const int out_off = - (((is_varlen ? 0 : seq_idx) * seq_len + local_t) * kNumVHeads + hv) * kHeadDimV + v_row; + (((is_varlen ? 0 : seq_idx) * seq_len + local_t) * kLocalVHeads + hv) * kHeadDimV + v_row; out[out_off] = cast_output(out_acc); } } @@ -195,7 +197,7 @@ struct Qwen35ScalarKdaPrefillKernel { } }; -template +template __global__ void qwen35_scalar_kda_prefill_kernel( const scalar_t* __restrict__ q, const scalar_t* __restrict__ k, @@ -213,8 +215,8 @@ __global__ void qwen35_scalar_kda_prefill_kernel( int sequence_count, bool is_varlen, bool has_initial_state) { - __shared__ typename Qwen35ScalarKdaPrefillKernel::SharedStorage storage; - Qwen35ScalarKdaPrefillKernel::run_device( + __shared__ typename Qwen35ScalarKdaPrefillKernel::SharedStorage storage; + Qwen35ScalarKdaPrefillKernel::run_device( q, k, v, @@ -234,7 +236,7 @@ __global__ void qwen35_scalar_kda_prefill_kernel( storage); } -template +template void launch_qwen35_scalar_kda_prefill_kernel( cudaStream_t stream, const scalar_t* q, @@ -253,9 +255,9 @@ void launch_qwen35_scalar_kda_prefill_kernel( int sequence_count, bool is_varlen, bool has_initial_state) { - const auto grid = Qwen35ScalarKdaPrefillKernel::grid_shape(sequence_count); - const auto block = Qwen35ScalarKdaPrefillKernel::block_shape(); - qwen35_scalar_kda_prefill_kernel<<>>( + const auto grid = Qwen35ScalarKdaPrefillKernel::grid_shape(sequence_count); + const auto block = Qwen35ScalarKdaPrefillKernel::block_shape(); + qwen35_scalar_kda_prefill_kernel<<>>( q, k, v, diff --git a/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu index 5af11afe..88b157b4 100644 --- a/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu +++ b/csrc/qwen35/prefill/sm90/qwen35_chunk_prefill_sm90.cu @@ -160,17 +160,17 @@ void qwen35_chunk_qk_prefill_sm90(at::Tensor q, at::Tensor k, at::Tensor out) { TORCH_CHECK(q.scalar_type() == at::kBFloat16, "q must be bfloat16"); TORCH_CHECK(k.scalar_type() == at::kBFloat16, "k must be bfloat16"); TORCH_CHECK(out.scalar_type() == at::kFloat, "out must be float32"); - TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B,T,48,128]"); - TORCH_CHECK(k.is_contiguous(), "k must be contiguous [B,T,48,128]"); - TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B,48,T,T]"); - TORCH_CHECK(q.dim() == 4, "q must be [B,T,48,128]"); + TORCH_CHECK(q.is_contiguous(), "q must be contiguous [B,T,HV,128]"); + TORCH_CHECK(k.is_contiguous(), "k must be contiguous [B,T,HV,128]"); + TORCH_CHECK(out.is_contiguous(), "out must be contiguous [B,HV,T,T]"); + TORCH_CHECK(q.dim() == 4, "q must be [B,T,HV,128]"); TORCH_CHECK(k.sizes() == q.sizes(), "k must match q"); const int64_t B = q.size(0); const int64_t T = q.size(1); const int64_t HV = q.size(2); - TORCH_CHECK(HV == kNumVHeads, "expected HV=48"); + TORCH_CHECK(decode::is_supported_local_v_heads(static_cast(HV)), "expected local HV in {48, 24, 12, 6}, got ", HV); TORCH_CHECK(q.size(3) == kHeadDimQK, "expected D=128"); - TORCH_CHECK(out.sizes() == at::IntArrayRef({B, HV, T, T}), "out must be [B,48,T,T]"); + TORCH_CHECK(out.sizes() == at::IntArrayRef({B, HV, T, T}), "out must be [B,HV,T,T]"); const at::cuda::OptionalCUDAGuard device_guard(q.device()); run_qwen35_chunk_qk_prefill_sm90_impl(q, k, out); diff --git a/cula/ops/qwen35_layout_decode.py b/cula/ops/qwen35_layout_decode.py index 1c63db41..eec5f975 100644 --- a/cula/ops/qwen35_layout_decode.py +++ b/cula/ops/qwen35_layout_decode.py @@ -77,6 +77,7 @@ def qwen35_layout_decode( if use_cudac: tokens = mixed_qkv_conv.shape[0] local_num_v_heads = a.shape[1] + infer_local_config(mixed_qkv_conv.shape[1], local_num_v_heads, config=config) q_rep = torch.empty( tokens, local_num_v_heads, diff --git a/cula/ops/qwen35_layout_prefill.py b/cula/ops/qwen35_layout_prefill.py index 64257d22..4a6a2867 100644 --- a/cula/ops/qwen35_layout_prefill.py +++ b/cula/ops/qwen35_layout_prefill.py @@ -73,8 +73,7 @@ def qwen35_layout_prefill( if use_cudac: tokens = mixed_qkv_conv.shape[0] local_num_v_heads = a.shape[1] - if local_num_v_heads != 48: - raise ValueError(f"backend='cudac' currently expects Qwen3.5 HV=48, got {local_num_v_heads}") + infer_local_config(mixed_qkv_conv.shape[1], local_num_v_heads, config=config) q_rep = torch.empty(tokens, local_num_v_heads, config.head_k_dim, device=mixed_qkv_conv.device, dtype=mixed_qkv_conv.dtype) k_rep = torch.empty_like(q_rep) v = torch.empty(tokens, local_num_v_heads, config.head_v_dim, device=mixed_qkv_conv.device, dtype=mixed_qkv_conv.dtype) diff --git a/cula/ops/qwen35_scalar_kda_prefill.py b/cula/ops/qwen35_scalar_kda_prefill.py index ae494356..a8ce4204 100644 --- a/cula/ops/qwen35_scalar_kda_prefill.py +++ b/cula/ops/qwen35_scalar_kda_prefill.py @@ -72,8 +72,8 @@ def qwen35_scalar_kda_prefill( raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_prefill is not available.") if use_cudac: - if HV != 48: - raise ValueError(f"backend='cudac' currently expects Qwen3.5 HV=48, got {HV}") + if HV not in (48, 24, 12, 6): + raise ValueError(f"backend='cudac' supports Qwen3.5 local HV in (48, 24, 12, 6), got {HV}") state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 out = torch.empty_like(v) final_state = torch.empty(state_count, HV, K, K, device=q.device, dtype=torch.float32) diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py index 1dbe50d4..1dc1f895 100644 --- a/tests/test_qwen35_decode.py +++ b/tests/test_qwen35_decode.py @@ -24,7 +24,7 @@ from cula.ops.qwen35_layout_decode import qwen35_layout_decode, qwen35_layout_decode_reference from cula.ops.qwen35_scalar_kda_decode import qwen35_layout_scalar_kda_decode, qwen35_scalar_kda_decode from cula.ops.qwen35_conv1d_decode import qwen35_conv1d_decode_reference, qwen35_conv1d_decode_update -from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig from cula.qwen35.runtime import qwen35_linear_attention_decode try: @@ -56,9 +56,13 @@ def _has_qwen35_fused_layout_kda_cudac(): return _has_qwen35_cudac() and hasattr(cula_cuda, "qwen35_layout_scalar_kda_decode") -def make_inputs(tokens: int = 2, pool_size: int = 3, device: torch.device | None = None): +def make_inputs( + tokens: int = 2, + pool_size: int = 3, + device: torch.device | None = None, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, +): device = _device() if device is None else device - config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG torch.manual_seed(0) mixed_qkv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) @@ -146,9 +150,10 @@ def manual_qwen35_layout_scalar_kda_reference( dt_bias: torch.Tensor, recurrent_state: torch.Tensor, state_indices: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, ): - config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG - q_rep, k_rep, v, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv_conv, a, b) + q_rep, k_rep, v, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv_conv, a, b, config=config) scale = config.head_k_dim**-0.5 q_f = torch.nn.functional.normalize(q_rep.float(), dim=-1) * scale @@ -179,6 +184,10 @@ def manual_qwen35_layout_scalar_kda_reference( return out.unsqueeze(1), state_out +def _local_config(local_v_heads: int) -> Qwen35LinearAttentionConfig: + return Qwen35LinearAttentionConfig(num_k_heads=local_v_heads // 3, num_v_heads=local_v_heads) + + @pytest.mark.parametrize("tokens", [1, 2]) def test_qwen35_conv_decode_reference(tokens: int): mixed_qkv, _, _, conv_weight, conv_state, _, _, _, _ = make_inputs(tokens=tokens) @@ -352,6 +361,49 @@ def test_qwen35_fused_layout_kda_cudac_matches_reference_unfused_and_triton(toke assert torch.allclose(state_triton, state_fused, atol=3e-5, rtol=3e-5) +@pytest.mark.skipif(not _has_qwen35_fused_layout_kda_cudac(), reason="Qwen3.5 fused layout+KDA CUDA backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_layout_scalar_kda_cudac_supports_local_tp_shards(local_v_heads: int): + config = _local_config(local_v_heads) + mixed_qkv, a, b, _, _, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=2, + pool_size=3, + device=torch.device("cuda"), + config=config, + ) + out_ref, state_ref = manual_qwen35_layout_scalar_kda_reference( + mixed_qkv, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices, + config=config, + ) + q_rep_ref, k_rep_ref, v_ref, a_ref, b_ref = qwen35_layout_decode_reference(mixed_qkv, a, b, config=config) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode(mixed_qkv, a, b, config=config, backend="cudac") + out, state = qwen35_layout_scalar_kda_decode( + mixed_qkv, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(q_rep, q_rep_ref) + torch.testing.assert_close(k_rep, k_rep_ref) + torch.testing.assert_close(v, v_ref) + torch.testing.assert_close(a_kernel, a_ref) + torch.testing.assert_close(b_kernel, b_ref) + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state, state_ref, atol=3e-5, rtol=3e-5) + + @pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") def test_qwen35_decode_cudac_rejects_duplicate_state_indices(): mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, _ = make_inputs( diff --git a/tests/test_qwen35_prefill.py b/tests/test_qwen35_prefill.py index d598b71a..445c9369 100644 --- a/tests/test_qwen35_prefill.py +++ b/tests/test_qwen35_prefill.py @@ -17,6 +17,7 @@ import sys import torch +import pytest sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) @@ -66,6 +67,10 @@ def run_seq(batch_idx, state_idx, start, end): return out, state +def _local_config(local_v_heads: int) -> Qwen35LinearAttentionConfig: + return Qwen35LinearAttentionConfig(num_k_heads=local_v_heads // 3, num_v_heads=local_v_heads) + + def test_qwen35_scalar_kda_prefill_reference_matches_manual(): torch.manual_seed(0) B, T, HV, K = 2, 3, 2, 128 @@ -170,6 +175,72 @@ def test_qwen35_scalar_kda_prefill_cuda_matches_reference(): torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_layout_prefill_cuda_supports_local_tp_shards(local_v_heads: int): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_layout_prefill"): + pytest.skip("qwen35_layout_prefill CUDA extension is not available") + + torch.manual_seed(20 + local_v_heads) + device = torch.device("cuda") + config = _local_config(local_v_heads) + tokens = 5 + mixed_qkv = torch.randn(tokens, config.conv_dim, device=device, dtype=torch.bfloat16) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=torch.bfloat16) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=torch.bfloat16) + + ref = qwen35_layout_prefill_reference(mixed_qkv, a, b, config=config) + out = qwen35_layout_prefill(mixed_qkv, a, b, config=config, backend="cudac") + + torch.cuda.synchronize() + for out_tensor, ref_tensor in zip(out, ref, strict=True): + torch.testing.assert_close(out_tensor, ref_tensor) + + +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_scalar_kda_prefill_cuda_supports_local_tp_shards(local_v_heads: int): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(30 + local_v_heads) + device = torch.device("cuda") + B, T, HV, K = 1, 4, local_v_heads, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="reference", + ) + out, state = qwen35_scalar_kda_prefill( + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state=initial_state, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + def test_qwen35_chunk_qk_prefill_sm90_matches_torch(): if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): import pytest @@ -190,6 +261,25 @@ def test_qwen35_chunk_qk_prefill_sm90_matches_torch(): torch.testing.assert_close(out, ref, atol=2e-1, rtol=2e-2) +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_chunk_qk_prefill_sm90_supports_local_tp_shards(local_v_heads: int): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): + pytest.skip("qwen35_chunk_qk_prefill_sm90 CUDA extension is not available") + + torch.manual_seed(40 + local_v_heads) + device = torch.device("cuda") + B, T, HV, K = 1, 32, local_v_heads, 128 + q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + out = torch.empty(B, HV, T, T, device=device, dtype=torch.float32) + + cula_cuda.qwen35_chunk_qk_prefill_sm90(q.contiguous(), k.contiguous(), out) + torch.cuda.synchronize() + + ref = torch.einsum("bthd,bshd->bhts", q.float(), k.float()) + torch.testing.assert_close(out, ref, atol=2e-1, rtol=2e-2) + + def test_qwen35_fused_kda_prefill_matches_reference(): if not torch.cuda.is_available(): import pytest From a7d4733106dfa98950a6cfcc3d2c0a7cfd214ec4 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Fri, 12 Jun 2026 06:33:45 +0000 Subject: [PATCH 11/35] Tune Qwen3.5 TP decode policies --- benchmarks/tune_qwen35_tp_policy.py | 362 ++++++++++++++++++ csrc/qwen35/decode/qwen35_decode_common.cuh | 27 ++ csrc/qwen35/decode/qwen35_layout_decode.cu | 9 +- csrc/qwen35/decode/qwen35_layout_kernel.hpp | 13 +- .../qwen35/decode/qwen35_scalar_kda_decode.cu | 8 +- .../decode/qwen35_scalar_kda_kernel.hpp | 20 +- cula/ops/qwen35_scalar_kda_decode.py | 5 +- cula/qwen35/runtime.py | 1 + tests/test_qwen35_decode.py | 181 ++++++++- 9 files changed, 600 insertions(+), 26 deletions(-) create mode 100644 benchmarks/tune_qwen35_tp_policy.py diff --git a/benchmarks/tune_qwen35_tp_policy.py b/benchmarks/tune_qwen35_tp_policy.py new file mode 100644 index 00000000..28653e9c --- /dev/null +++ b/benchmarks/tune_qwen35_tp_policy.py @@ -0,0 +1,362 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tune Qwen3.5 TP-local kernel policies. + +This is a configuration-driven tuner. It benchmarks only policies that are +compiled into the current extension and records unsupported candidates in the +result file. The initial compiled policy is the decode traits currently used by +the CUDA/CuTe kernels: + + layout_vec=4, kda_threads=128, kda_tile_v=16, kda_tile_k=16, heads_per_cta=1 + +When more C++ policy specializations are added, extend `compiled_policy_key` +and the kernel dispatch path; this script can then sweep them without changing +the output format. +""" + +from __future__ import annotations + +import argparse +import csv +import itertools +import json +import pathlib +import sys +from dataclasses import asdict, dataclass +from typing import Any + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + + +@dataclass(frozen=True) +class DecodePolicy: + name: str + layout_vec: int + kda_threads: int + kda_tile_v: int + kda_tile_k: int + heads_per_cta: int = 1 + + @property + def key(self) -> tuple[int, int, int, int, int]: + return (self.layout_vec, self.kda_threads, self.kda_tile_v, self.kda_tile_k, self.heads_per_cta) + + +CURRENT_DECODE_POLICY = DecodePolicy( + name="current", + layout_vec=4, + kda_threads=128, + kda_tile_v=16, + kda_tile_k=16, + heads_per_cta=1, +) + + +def decode_benchmarks(): + from benchmarks import bench_qwen35_decode + + return bench_qwen35_decode + + +def compiled_policy_key(policy: DecodePolicy) -> str | None: + """Return the compiled backend selector for a policy, or None if absent.""" + if policy.key == CURRENT_DECODE_POLICY.key: + return "current" + return None + + +def _list_from_json(data: dict[str, Any], key: str, default: list[int]) -> list[int]: + value = data.get(key, default) + if not isinstance(value, list) or not value: + raise ValueError(f"{key} must be a non-empty list") + return [int(item) for item in value] + + +def load_decode_policies(path: pathlib.Path | None) -> list[DecodePolicy]: + if path is None: + return [CURRENT_DECODE_POLICY] + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + + if isinstance(data, list): + policies = [] + for idx, item in enumerate(data): + if not isinstance(item, dict): + raise ValueError(f"Policy entry {idx} must be an object") + policies.append( + DecodePolicy( + name=str(item.get("name", f"policy_{idx}")), + layout_vec=int(item["layout_vec"]), + kda_threads=int(item["kda_threads"]), + kda_tile_v=int(item["kda_tile_v"]), + kda_tile_k=int(item["kda_tile_k"]), + heads_per_cta=int(item.get("heads_per_cta", 1)), + ) + ) + return policies + + if not isinstance(data, dict): + raise ValueError("Policy grid must be a JSON object or list") + if data.get("mode", "decode") != "decode": + raise ValueError("Only decode policy grids are supported by this tuner") + + policies = [] + for idx, combo in enumerate( + itertools.product( + _list_from_json(data, "layout_vec", [CURRENT_DECODE_POLICY.layout_vec]), + _list_from_json(data, "kda_threads", [CURRENT_DECODE_POLICY.kda_threads]), + _list_from_json(data, "kda_tile_v", [CURRENT_DECODE_POLICY.kda_tile_v]), + _list_from_json(data, "kda_tile_k", [CURRENT_DECODE_POLICY.kda_tile_k]), + _list_from_json(data, "heads_per_cta", [CURRENT_DECODE_POLICY.heads_per_cta]), + ) + ): + layout_vec, kda_threads, kda_tile_v, kda_tile_k, heads_per_cta = combo + policies.append( + DecodePolicy( + name=f"p{idx}_lv{layout_vec}_th{kda_threads}_tv{kda_tile_v}_tk{kda_tile_k}_h{heads_per_cta}", + layout_vec=layout_vec, + kda_threads=kda_threads, + kda_tile_v=kda_tile_v, + kda_tile_k=kda_tile_k, + heads_per_cta=heads_per_cta, + ) + ) + return policies + + +def write_example_grid(path: pathlib.Path) -> None: + example = { + "mode": "decode", + "layout_vec": [4, 8], + "kda_threads": [64, 128, 256], + "kda_tile_v": [8, 16, 32], + "kda_tile_k": [8, 16, 32], + "heads_per_cta": [1, 2, 4], + } + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(example, f, indent=2, sort_keys=True) + + +def bucket_name(tokens: int) -> str: + if tokens <= 4: + return "tokens<=4" + if tokens <= 16: + return "tokens<=16" + if tokens <= 64: + return "tokens<=64" + return "tokens>64" + + +def run_decode_policy( + *, + scope: str, + tokens: int, + tp_size: int, + warmup: int, + rep: int, + seed: int, + policy: DecodePolicy, +) -> dict[str, Any]: + decode_bench = decode_benchmarks() + config = decode_bench.local_config_from_tp_size(tp_size) + compiled_key = compiled_policy_key(policy) + row: dict[str, Any] = { + "mode": "decode", + "scope": scope, + "tokens": tokens, + "token_bucket": bucket_name(tokens), + "tp_size": tp_size, + "local_k_heads": config.num_k_heads, + "local_v_heads": config.num_v_heads, + "conv_dim": config.conv_dim, + "policy": policy.name, + "compiled_policy": compiled_key, + **asdict(policy), + } + if compiled_key is None: + row.update({"status": "unsupported", "ms": None, "us_per_token": None}) + return row + + device = decode_bench.accelerator_device() + if scope == "core": + ms = decode_bench.bench_native_core(tokens, device, warmup, rep, seed, config) + elif scope == "fused": + ms = decode_bench.bench_fused_layout_kda(tokens, device, warmup, rep, seed, config) + elif scope == "full": + ms = decode_bench.bench_full(tokens, device, warmup, rep, seed, config) + else: + raise ValueError(f"Unsupported decode scope={scope}") + + row.update({"status": "ok", "ms": ms, "us_per_token": ms * 1000.0 / tokens}) + return row + + +def choose_best(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: + groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} + for row in rows: + if row["status"] != "ok": + continue + key = (row["mode"], row["scope"], row["tp_size"], row["local_v_heads"], row["token_bucket"]) + groups.setdefault(key, []).append(row) + + best_rows = [] + for key, candidates in sorted(groups.items()): + best = min(candidates, key=lambda row: float(row["ms"])) + mode, scope, tp_size, local_v_heads, token_bucket = key + best_rows.append( + { + "mode": mode, + "scope": scope, + "tp_size": tp_size, + "local_v_heads": local_v_heads, + "token_bucket": token_bucket, + "policy": best["policy"], + "compiled_policy": best["compiled_policy"], + "ms": best["ms"], + "us_per_token": best["us_per_token"], + "layout_vec": best["layout_vec"], + "kda_threads": best["kda_threads"], + "kda_tile_v": best["kda_tile_v"], + "kda_tile_k": best["kda_tile_k"], + "heads_per_cta": best["heads_per_cta"], + } + ) + return best_rows + + +def write_csv(path: pathlib.Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fieldnames = [ + "mode", + "scope", + "tokens", + "token_bucket", + "tp_size", + "local_k_heads", + "local_v_heads", + "conv_dim", + "policy", + "compiled_policy", + "status", + "ms", + "us_per_token", + "layout_vec", + "kda_threads", + "kda_tile_v", + "kda_tile_k", + "heads_per_cta", + ] + with path.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + writer.writerows(rows) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Tune Qwen3.5 TP-local kernel policies.") + parser.add_argument("--mode", choices=["decode"], default="decode") + parser.add_argument("--scope", choices=["core", "fused", "full", "all"], default="fused") + parser.add_argument("--tp-sizes", nargs="+", type=int, choices=[1, 2, 4, 8], default=[1, 2, 4, 8]) + parser.add_argument("--tokens", nargs="+", type=int, default=[1, 2, 4, 8, 16, 32, 64, 128]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--policy-grid", type=pathlib.Path, default=None) + parser.add_argument("--write-example-grid", type=pathlib.Path, default=None) + parser.add_argument("--output-json", type=pathlib.Path, default=pathlib.Path("tmp/qwen35_tp_policy_tune.json")) + parser.add_argument("--csv", type=pathlib.Path, default=None) + parser.add_argument("--fail-on-unsupported", action="store_true") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.write_example_grid is not None: + write_example_grid(args.write_example_grid) + print(f"wrote example policy grid: {args.write_example_grid}") + return 0 + + policies = load_decode_policies(args.policy_grid) + scopes = ["core", "fused", "full"] if args.scope == "all" else [args.scope] + decode_bench = decode_benchmarks() + device = decode_bench.accelerator_device() + device_name = decode_bench.accelerator_name(device) + + print(f"Qwen3.5 TP policy tuner: mode={args.mode} device={device_name}") + print(f"tp_sizes={args.tp_sizes} tokens={args.tokens} scopes={scopes}") + print(f"policies={len(policies)} compiled={sum(compiled_policy_key(p) is not None for p in policies)}") + + rows: list[dict[str, Any]] = [] + for policy in policies: + compiled_key = compiled_policy_key(policy) + if compiled_key is None: + print(f"skip unsupported policy={policy.name} {asdict(policy)}") + for scope in scopes: + for tp_size in args.tp_sizes: + for tokens in args.tokens: + row = run_decode_policy( + scope=scope, + tokens=tokens, + tp_size=tp_size, + warmup=args.warmup, + rep=args.rep, + seed=args.seed, + policy=policy, + ) + rows.append(row) + if row["status"] == "ok": + print( + f"{scope:>5} tp={tp_size} hv={row['local_v_heads']:>2} tokens={tokens:>4} " + f"policy={policy.name} ms={row['ms']:.4f} us/tok={row['us_per_token']:.2f}" + ) + + unsupported = [row for row in rows if row["status"] == "unsupported"] + if unsupported and args.fail_on_unsupported: + raise RuntimeError(f"{len(unsupported)} policy/shape rows are unsupported by the compiled extension") + + best_rows = choose_best(rows) + result = { + "device": device_name, + "mode": args.mode, + "warmup": args.warmup, + "rep": args.rep, + "seed": args.seed, + "rows": rows, + "best": best_rows, + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + with args.output_json.open("w", encoding="utf-8") as f: + json.dump(result, f, indent=2, sort_keys=True) + print(f"wrote {args.output_json}") + + if args.csv is not None: + write_csv(args.csv, rows) + print(f"wrote {args.csv}") + + if best_rows: + print("best policies:") + for row in best_rows: + print( + f" {row['scope']:>5} tp={row['tp_size']} hv={row['local_v_heads']:>2} " + f"{row['token_bucket']}: {row['policy']} {row['ms']:.4f} ms" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/csrc/qwen35/decode/qwen35_decode_common.cuh b/csrc/qwen35/decode/qwen35_decode_common.cuh index 061691cb..51ce1872 100644 --- a/csrc/qwen35/decode/qwen35_decode_common.cuh +++ b/csrc/qwen35/decode/qwen35_decode_common.cuh @@ -50,6 +50,33 @@ inline constexpr bool is_supported_local_v_heads(int local_v_heads) { return local_v_heads == 48 || local_v_heads == 24 || local_v_heads == 12 || local_v_heads == 6; } +template +struct Qwen35DecodeLocalShape { + static_assert(is_supported_local_v_heads(kLocalVHeads_), "Unsupported Qwen3.5 local V-head count."); + static constexpr int kLocalVHeads = kLocalVHeads_; + static constexpr int kLocalQKHeads = local_qk_heads_from_v_heads(kLocalVHeads); + static constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; + static constexpr int kLocalQDim = local_q_dim(kLocalQKHeads); + static constexpr int kLocalKDim = local_q_dim(kLocalQKHeads); + static constexpr int kLocalVDim = local_v_dim(kLocalVHeads); + static constexpr int kLocalMixedQKVDim = local_mixed_qkv_dim(kLocalQKHeads, kLocalVHeads); + + // Decode shape policy. Head dimension is fixed at 128 for Qwen3.5, but keep + // these knobs with the local-head traits so future TP-shape tuning has one + // place to specialize. + static constexpr int kLayoutVec = 4; + static constexpr int kLayoutThreads = kHeadDimQK / kLayoutVec; + static constexpr int kKdaThreads = 128; + static constexpr int kKdaTileV = 16; + static constexpr int kKdaTileK = 16; + + static_assert(kLocalVHeads % kLocalQKHeads == 0); + static_assert(kHeadDimQK == kHeadDimV); + static_assert(kHeadDimQK % kLayoutVec == 0); + static_assert(kHeadDimV % kKdaTileV == 0); + static_assert(kHeadDimQK % kKdaTileK == 0); +}; + struct ConvDecodeParams { at::Tensor mixed_qkv; // [B, 1, local_conv_dim] at::Tensor conv_state; // [B, local_conv_dim, 4] diff --git a/csrc/qwen35/decode/qwen35_layout_decode.cu b/csrc/qwen35/decode/qwen35_layout_decode.cu index 138730ea..811e98b9 100644 --- a/csrc/qwen35/decode/qwen35_layout_decode.cu +++ b/csrc/qwen35/decode/qwen35_layout_decode.cu @@ -49,11 +49,10 @@ void launch_layout_decode_for_heads( scalar_t* a_kernel, scalar_t* b_kernel, int64_t batch_size) { - constexpr int kLocalQKHeads = cula::qwen35::decode::local_qk_heads_from_v_heads(kLocalVHeads); - constexpr int threads = 32; - dim3 grid(kLocalVHeads, static_cast(batch_size), 1); - cula::qwen35::decode::qwen35_layout_decode_kernel_cute - <<>>( + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + dim3 grid(Shape::kLocalVHeads, static_cast(batch_size), 1); + cula::qwen35::decode::qwen35_layout_decode_kernel_cute + <<>>( mixed_qkv_conv, a, b, diff --git a/csrc/qwen35/decode/qwen35_layout_kernel.hpp b/csrc/qwen35/decode/qwen35_layout_kernel.hpp index b665dce2..e43ffc00 100644 --- a/csrc/qwen35/decode/qwen35_layout_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_layout_kernel.hpp @@ -56,17 +56,18 @@ __global__ void qwen35_layout_decode_kernel_cute( scalar_t* __restrict__ a_kernel, scalar_t* __restrict__ b_kernel, int64_t token_count) { - static_assert(kLocalVHeads % kLocalQKHeads == 0); - constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; - constexpr int kLocalQDim = kLocalQKHeads * kHeadDimQK; - constexpr int kLocalKDim = kLocalQKHeads * kHeadDimQK; - constexpr int kLocalMixedQKVDim = 2 * kLocalQDim + kLocalVHeads * kHeadDimV; + using Shape = Qwen35DecodeLocalShape; + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; // TODO(qwen35-layout-opt): // - Re-evaluate whether Vec=8 is profitable for bf16/fp16 on the target GPUs. // - Push more of the q/k repeat mapping into compile-time CuTe layout transforms. // - Revisit whether a shared-memory staging path is worthwhile after profiling. // - Consider widening the a/b writeback path if it shows up in profiling. - constexpr int kVec = 4; + constexpr int kVec = Shape::kLayoutVec; static_assert(kHeadDimV % kVec == 0); static_assert(kHeadDimQK == kHeadDimV); static_assert(kHeadDimQK % kVec == 0); diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu index 8718b4ff..ab822752 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu +++ b/csrc/qwen35/decode/qwen35_scalar_kda_decode.cu @@ -42,8 +42,8 @@ void dispatch_scalar_decode_for_heads( const int32_t* pool_idx, scalar_t* out, int token_count) { - constexpr int kLocalQKHeads = local_qk_heads_from_v_heads(kLocalVHeads); - kernel::launch_qwen35_scalar_kda_decode_kernel( + using Shape = Qwen35DecodeLocalShape; + kernel::launch_qwen35_scalar_kda_decode_kernel( stream, q_rep, k_rep, @@ -70,8 +70,8 @@ void dispatch_layout_scalar_decode_for_heads( const int32_t* pool_idx, scalar_t* out, int token_count) { - constexpr int kLocalQKHeads = local_qk_heads_from_v_heads(kLocalVHeads); - kernel::launch_qwen35_layout_scalar_kda_decode_kernel( + using Shape = Qwen35DecodeLocalShape; + kernel::launch_qwen35_layout_scalar_kda_decode_kernel( stream, mixed_qkv_conv, a, diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 3cd032f4..9c91e8b3 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -25,6 +25,8 @@ using namespace cute; template struct Qwen35ScalarKdaDecodeKernel { + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); // Decode-first design: // - 1 CTA owns 1 (token_idx, hv) // - 1 warpgroup (128 threads) per CTA @@ -32,10 +34,10 @@ struct Qwen35ScalarKdaDecodeKernel { // internal [V, K] view // - the intended optimized path is fp32 FFMA on CUDA cores, not a forced // Tensor Core lowering - static constexpr int kThreads = 128; - static constexpr int kWarpGroupThreads = 128; - static constexpr int kTileV = 16; - static constexpr int kTileK = 16; + static constexpr int kThreads = Shape::kKdaThreads; + static constexpr int kWarpGroupThreads = Shape::kKdaThreads; + static constexpr int kTileV = Shape::kKdaTileV; + static constexpr int kTileK = Shape::kKdaTileK; static constexpr int kTilesPerV = kHeadDimV / kTileV; static constexpr int kTilesPerK = kHeadDimQK / kTileK; @@ -65,7 +67,7 @@ struct Qwen35ScalarKdaDecodeKernel { static dim3 grid_shape(int token_count) { // One block owns one (token_idx, hv) pair in the first implementation. - return dim3(static_cast(kLocalVHeads), static_cast(token_count), 1); + return dim3(static_cast(Shape::kLocalVHeads), static_cast(token_count), 1); } template @@ -197,10 +199,10 @@ struct Qwen35ScalarKdaDecodeKernel { scalar_t* __restrict__ out, int token_count, SharedStorage& storage) { - constexpr int kRepeatFactor = kLocalVHeads / kLocalQKHeads; - constexpr int kLocalQDim = kLocalQKHeads * kHeadDimQK; - constexpr int kLocalKDim = kLocalQKHeads * kHeadDimQK; - constexpr int kLocalMixedQKVDim = 2 * kLocalQDim + kLocalVHeads * kHeadDimV; + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; const int hv = static_cast(blockIdx.x); const int token_idx = static_cast(blockIdx.y); diff --git a/cula/ops/qwen35_scalar_kda_decode.py b/cula/ops/qwen35_scalar_kda_decode.py index c32c30ee..45d2c6a6 100644 --- a/cula/ops/qwen35_scalar_kda_decode.py +++ b/cula/ops/qwen35_scalar_kda_decode.py @@ -18,6 +18,8 @@ import torch +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, Qwen35LinearAttentionConfig + try: import cula.cudac as cula_cuda except ImportError: @@ -146,6 +148,7 @@ def qwen35_layout_scalar_kda_decode( recurrent_state: torch.Tensor, *, state_indices: torch.Tensor | None = None, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, backend: str = "auto", ) -> tuple[torch.Tensor, torch.Tensor]: """Fused Qwen3.5 layout decode + scalar-gated KDA decode.""" @@ -209,7 +212,7 @@ def qwen35_layout_scalar_kda_decode( from cula.ops.qwen35_layout_decode import qwen35_layout_decode_reference - q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode_reference(mixed_qkv_conv, a, b) + q_rep, k_rep, v, a_kernel, b_kernel = qwen35_layout_decode_reference(mixed_qkv_conv, a, b, config=config) return qwen35_scalar_kda_decode( q=q_rep.unsqueeze(1).contiguous(), k=k_rep.unsqueeze(1).contiguous(), diff --git a/cula/qwen35/runtime.py b/cula/qwen35/runtime.py index 6be6371e..f3f55091 100644 --- a/cula/qwen35/runtime.py +++ b/cula/qwen35/runtime.py @@ -335,6 +335,7 @@ def qwen35_linear_attention_decode( dt_bias=dt_bias, recurrent_state=recurrent_state, state_indices=state_indices, + config=config, backend=backend, ) core_attn_out = core_attn_out.reshape(tokens, local_value_dim) diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py index 1dc1f895..274ff85d 100644 --- a/tests/test_qwen35_decode.py +++ b/tests/test_qwen35_decode.py @@ -106,8 +106,9 @@ def manual_qwen35_decode_reference( conv_state: torch.Tensor, recurrent_state: torch.Tensor, state_indices: torch.Tensor, + *, + config: Qwen35LinearAttentionConfig = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, ): - config = DEFAULT_QWEN35_LINEAR_ATTN_CONFIG conv_out, conv_state_out = manual_conv_decode(mixed_qkv, conv_state, conv_weight) q_end = config.key_dim k_end = q_end + config.key_dim @@ -142,6 +143,45 @@ def manual_qwen35_decode_reference( return out, conv_state_out, state_out +def manual_qwen35_scalar_kda_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + recurrent_state: torch.Tensor, + state_indices: torch.Tensor, +): + if a.ndim == 2: + a = a.unsqueeze(1) + if b.ndim == 2: + b = b.unsqueeze(1) + N, _, HV, K = q.shape + scale = K**-0.5 + q_f = torch.nn.functional.normalize(q.squeeze(1).float(), dim=-1) * scale + k_f = torch.nn.functional.normalize(k.squeeze(1).float(), dim=-1) + v_f = v.squeeze(1).float() + state_out = recurrent_state.clone() + out = torch.empty(N, 1, HV, v.shape[-1], device=q.device, dtype=v.dtype) + + for token_idx in range(N): + pool_idx = int(state_indices[token_idx].item()) + for hv in range(HV): + state_kv = state_out[pool_idx, hv] + decay = torch.exp(-torch.exp(A_log[hv]) * torch.nn.functional.softplus(a[token_idx, 0, hv].float() + dt_bias[hv])) + beta = torch.sigmoid(b[token_idx, 0, hv].float()) + k_vec = k_f[token_idx, hv] + q_vec = q_f[token_idx, hv] + proj = decay * (state_kv.transpose(0, 1) @ k_vec) + v_new = beta * (v_f[token_idx, hv] - proj) + state_new_kv = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) + out[token_idx, 0, hv] = (state_new_kv.transpose(0, 1) @ q_vec).to(v.dtype) + state_out[pool_idx, hv] = state_new_kv + return out, state_out + + def manual_qwen35_layout_scalar_kda_reference( mixed_qkv_conv: torch.Tensor, a: torch.Tensor, @@ -284,6 +324,130 @@ def test_qwen35_decode_cudac_matches_reference(tokens: int): assert torch.allclose(recurrent_state_ref, recurrent_state_out, atol=3e-5, rtol=3e-5) +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_conv_decode_cudac_supports_local_tp_shapes(local_v_heads: int): + config = _local_config(local_v_heads) + mixed_qkv, _, _, conv_weight, conv_state, _, _, _, _ = make_inputs( + tokens=3, + pool_size=3, + device=torch.device("cuda"), + config=config, + ) + y_ref, state_ref = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + backend="reference", + ) + y, state = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(y, y_ref) + torch.testing.assert_close(state, state_ref) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_scalar_kda_decode_cudac_supports_local_tp_shapes(local_v_heads: int): + torch.manual_seed(3) + config = _local_config(local_v_heads) + tokens = 3 + device = torch.device("cuda") + q = torch.randn(tokens, 1, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + k = torch.randn_like(q) + v = torch.randn(tokens, 1, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, 1, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, 1, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + recurrent_state = torch.randn( + tokens, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + + out_ref, state_ref = manual_qwen35_scalar_kda_reference( + q, + k, + v, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices, + ) + out, state = qwen35_scalar_kda_decode( + q, + k, + v, + a, + b, + A_log, + dt_bias, + recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state, state_ref, atol=3e-5, rtol=3e-5) + + +@pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") +@pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) +def test_qwen35_decode_cudac_supports_local_tp_shapes(local_v_heads: int): + config = _local_config(local_v_heads) + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=3, + pool_size=3, + device=torch.device("cuda"), + config=config, + ) + out_ref, conv_state_ref, recurrent_state_ref = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="reference", + ) + out, conv_state_out, recurrent_state_out = qwen35_linear_attention_decode( + mixed_qkv, + a, + b, + conv_weight, + A_log, + dt_bias, + config=config, + conv_state=conv_state, + recurrent_state=recurrent_state, + state_indices=state_indices, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(conv_state_out, conv_state_ref) + torch.testing.assert_close(recurrent_state_out, recurrent_state_ref, atol=3e-5, rtol=3e-5) + + @pytest.mark.skipif(not _has_qwen35_fused_layout_kda_cudac(), reason="Qwen3.5 fused layout+KDA CUDA backend is not available") @pytest.mark.parametrize("tokens", [1, 2, 4]) def test_qwen35_fused_layout_kda_cudac_matches_reference_unfused_and_triton(tokens: int): @@ -348,6 +512,7 @@ def test_qwen35_fused_layout_kda_cudac_matches_reference_unfused_and_triton(toke dt_bias=dt_bias, recurrent_state=recurrent_state, state_indices=state_indices, + config=DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, backend="cudac", ) @@ -391,6 +556,18 @@ def test_qwen35_layout_scalar_kda_cudac_supports_local_tp_shards(local_v_heads: dt_bias, recurrent_state, state_indices=state_indices, + config=config, + backend="cudac", + ) + out_3d_gate, state_3d_gate = qwen35_layout_scalar_kda_decode( + mixed_qkv, + a.unsqueeze(1), + b.unsqueeze(1), + A_log, + dt_bias, + recurrent_state, + state_indices=state_indices, + config=config, backend="cudac", ) @@ -402,6 +579,8 @@ def test_qwen35_layout_scalar_kda_cudac_supports_local_tp_shards(local_v_heads: torch.testing.assert_close(b_kernel, b_ref) torch.testing.assert_close(out.float(), out_ref.float(), atol=3e-2, rtol=3e-2) torch.testing.assert_close(state, state_ref, atol=3e-5, rtol=3e-5) + torch.testing.assert_close(out_3d_gate, out) + torch.testing.assert_close(state_3d_gate, state) @pytest.mark.skipif(not _has_qwen35_cudac(), reason="Qwen3.5 CUDA decode backend is not available") From b38e9a2df1f69ad9fd9ca8512026a5e737d2f7ac Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Fri, 12 Jun 2026 10:30:56 +0000 Subject: [PATCH 12/35] Optimize qwen35 decode mainloop tiling --- csrc/qwen35/decode/qwen35_decode_common.cuh | 4 +- .../decode/qwen35_scalar_kda_kernel.hpp | 21 ++-- .../decode/qwen35_scalar_kda_mainloop.hpp | 97 ++++++++++++------- 3 files changed, 80 insertions(+), 42 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_decode_common.cuh b/csrc/qwen35/decode/qwen35_decode_common.cuh index 51ce1872..2f78d22e 100644 --- a/csrc/qwen35/decode/qwen35_decode_common.cuh +++ b/csrc/qwen35/decode/qwen35_decode_common.cuh @@ -67,8 +67,8 @@ struct Qwen35DecodeLocalShape { static constexpr int kLayoutVec = 4; static constexpr int kLayoutThreads = kHeadDimQK / kLayoutVec; static constexpr int kKdaThreads = 128; - static constexpr int kKdaTileV = 16; - static constexpr int kKdaTileK = 16; + static constexpr int kKdaTileV = 32; + static constexpr int kKdaTileK = kHeadDimQK; static_assert(kLocalVHeads % kLocalQKHeads == 0); static_assert(kHeadDimQK == kHeadDimV); diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 9c91e8b3..ed67ccc6 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -57,6 +57,7 @@ struct Qwen35ScalarKdaDecodeKernel { alignas(16) float k_smem[kHeadDimQK]; alignas(16) scalar_t v_smem[kHeadDimV]; alignas(16) float norm_smem[2]; + alignas(16) float state_smem[kHeadDimQK * kTileV]; alignas(16) float proj_smem[kHeadDimV]; alignas(16) float out_smem[kHeadDimV]; }; @@ -66,8 +67,12 @@ struct Qwen35ScalarKdaDecodeKernel { } static dim3 grid_shape(int token_count) { - // One block owns one (token_idx, hv) pair in the first implementation. - return dim3(static_cast(Shape::kLocalVHeads), static_cast(token_count), 1); + // One block owns one V tile for one (token_idx, hv) pair. + constexpr int kNumVTiles = (kHeadDimV + kTileV - 1) / kTileV; + return dim3( + static_cast(kNumVTiles), + static_cast(Shape::kLocalVHeads), + static_cast(token_count)); } template @@ -84,8 +89,9 @@ struct Qwen35ScalarKdaDecodeKernel { scalar_t* __restrict__ out, int token_count, SharedStorage& storage) { - const int hv = static_cast(blockIdx.x); - const int token_idx = static_cast(blockIdx.y); + const int v_tile = static_cast(blockIdx.x); + const int hv = static_cast(blockIdx.y); + const int token_idx = static_cast(blockIdx.z); const int tid = static_cast(threadIdx.x); if (token_idx >= token_count || hv >= kLocalVHeads) { return; @@ -183,6 +189,7 @@ struct Qwen35ScalarKdaDecodeKernel { state_vk, out_vec, storage, + v_tile * kTileV, tid, kThreads); } @@ -204,8 +211,9 @@ struct Qwen35ScalarKdaDecodeKernel { constexpr int kLocalKDim = Shape::kLocalKDim; constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; - const int hv = static_cast(blockIdx.x); - const int token_idx = static_cast(blockIdx.y); + const int v_tile = static_cast(blockIdx.x); + const int hv = static_cast(blockIdx.y); + const int token_idx = static_cast(blockIdx.z); const int tid = static_cast(threadIdx.x); if (token_idx >= token_count || hv >= kLocalVHeads) { return; @@ -275,6 +283,7 @@ struct Qwen35ScalarKdaDecodeKernel { state_vk, out_vec, storage, + v_tile * kTileV, tid, kThreads); } diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp index 0c9401bc..206cbf91 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp @@ -56,8 +56,8 @@ struct Qwen35ScalarKdaDecodeMainloop { // 2. optimize proj / update / out reductions // 3. evaluate warp-specialized load/compute roles only after the fp32 path // is stable and measured - static constexpr int kTileV = 16; - static constexpr int kTileK = 16; + static constexpr int kTileV = 32; + static constexpr int kTileK = kHeadDimQK; static constexpr int kTilesPerV = kHeadDimV / kTileV; static constexpr int kTilesPerK = kHeadDimQK / kTileK; static constexpr int kWarpSize = 32; @@ -65,10 +65,14 @@ struct Qwen35ScalarKdaDecodeMainloop { static constexpr int kWarpsPerCta = 4; static constexpr int kRowsPerWarp = kWarpSize; static constexpr int kRowsPerThread = 1; + static constexpr int kThreadsPerVRow = 4; + static constexpr int kKPerThread = kHeadDimQK / kThreadsPerVRow; static_assert(kHeadDimV == 128); static_assert(kHeadDimQK == 128); static_assert(kWarpsPerCta * kRowsPerWarp == kHeadDimV); + static_assert(kTileV * kThreadsPerVRow == kWarpsPerCta * kWarpSize); + static_assert(kHeadDimQK % kThreadsPerVRow == 0); // First concrete decode threading plan: // @@ -358,6 +362,7 @@ struct Qwen35ScalarKdaDecodeMainloop { TensorHvk& state_vk, TensorOut& out_vec, SharedStorage& storage, + int v_tile_base, int tid, int num_threads) { // Decode organization: @@ -393,6 +398,9 @@ struct Qwen35ScalarKdaDecodeMainloop { auto k_smem = make_tensor(make_smem_ptr(storage.k_smem), make_layout(make_shape(Int{}))); auto v_smem = make_tensor(make_smem_ptr(storage.v_smem), make_layout(make_shape(Int{}))); auto norm_smem = make_tensor(make_smem_ptr(storage.norm_smem), make_layout(make_shape(Int<2>{}))); + auto state_smem = make_tensor( + make_smem_ptr(storage.state_smem), + make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); auto proj_smem = make_tensor(make_smem_ptr(storage.proj_smem), make_layout(make_shape(Int{}))); auto out_smem = make_tensor(make_smem_ptr(storage.out_smem), make_layout(make_shape(Int{}))); @@ -401,11 +409,22 @@ struct Qwen35ScalarKdaDecodeMainloop { q_smem(idx) = static_cast(q_vec(idx)); k_smem(idx) = static_cast(k_vec(idx)); } + for (int local_v = tid; local_v < kTileV; local_v += num_threads) { + const int v_global = v_tile_base + local_v; + if (v_global < kHeadDimV) { + v_smem(local_v) = v_vec(v_global); + } + } for (int idx = tid; idx < kHeadDimV; idx += num_threads) { - v_smem(idx) = v_vec(idx); proj_smem(idx) = 0.f; out_smem(idx) = 0.f; } + for (int idx = tid; idx < kHeadDimQK * kTileV; idx += num_threads) { + const int k_idx = idx / kTileV; + const int local_v = idx - k_idx * kTileV; + const int v_global = v_tile_base + local_v; + state_smem(k_idx, local_v) = v_global < kHeadDimV ? static_cast(state_vk(v_global, k_idx)) : 0.f; + } __syncthreads(); if (tid == 0) { @@ -429,45 +448,55 @@ struct Qwen35ScalarKdaDecodeMainloop { } __syncthreads(); - ThreadRowPlan row_plan = make_thread_row_plan(tid); + const int local_v = tid / kThreadsPerVRow; + const int row_lane = tid - local_v * kThreadsPerVRow; + const int v_global = v_tile_base + local_v; + if (local_v < kTileV && v_global < kHeadDimV) { + const int k_begin = row_lane * kKPerThread; + float proj_part = 0.f; +#pragma unroll + for (int kk = 0; kk < kKPerThread; ++kk) { + const int k_idx = k_begin + kk; + proj_part += state_smem(k_idx, local_v) * k_smem(k_idx); + } + proj_smem(tid) = proj_part; + } + __syncthreads(); - // First concrete ownership model: - // - each thread owns one full state row across all 128 K columns - // - the row is streamed tile-by-tile through registers - // - no cross-thread reduction is needed for proj/out because the full row - // stays with one thread for the duration of the token update - if (row_plan.owns_row) { + if (local_v < kTileV && row_lane == 0 && v_global < kHeadDimV) { float proj_row = 0.f; - for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { - TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; - RowTileProjPlan proj_plan = make_row_tile_proj_plan(coords, row_plan.warp_id, row_plan.lane_id); - proj_plan.v_row = row_plan.v_row; - proj_plan.owns_row = true; - proj_row += project_row_tile(state_vk, k_smem, proj_plan); +#pragma unroll + for (int lane = 0; lane < kThreadsPerVRow; ++lane) { + proj_row += proj_smem(local_v * kThreadsPerVRow + lane); } + const float v_val = static_cast(v_smem(local_v)); + proj_smem(local_v) = beta * (v_val - decay * proj_row); + } + __syncthreads(); - proj_smem(row_plan.v_row) = proj_row; - - const float v_val = static_cast(v_smem(row_plan.v_row)); - const float decayed_proj_row = decay * proj_row; - const float v_new_row = beta * (v_val - decayed_proj_row); - - float out_row = 0.f; - for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { - TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; - RowTileUpdatePlan update_plan = make_row_tile_update_plan(coords, row_plan.warp_id, row_plan.lane_id); - update_plan.v_row = row_plan.v_row; - update_plan.owns_row = true; - out_row += update_and_output_row_tile( - state_vk, k_smem, q_smem, update_plan, decay, v_new_row); + if (local_v < kTileV && v_global < kHeadDimV) { + const int k_begin = row_lane * kKPerThread; + const float v_new_row = proj_smem(local_v); + float out_part = 0.f; +#pragma unroll + for (int kk = 0; kk < kKPerThread; ++kk) { + const int k_idx = k_begin + kk; + const float state_new = decay * state_smem(k_idx, local_v) + v_new_row * k_smem(k_idx); + state_smem(k_idx, local_v) = state_new; + out_part += state_new * q_smem(k_idx); + state_vk(v_global, k_idx) = state_new; } - - out_smem(row_plan.v_row) = out_row; + out_smem(local_v * kThreadsPerVRow + row_lane) = out_part; } __syncthreads(); - for (int idx = tid; idx < kHeadDimV; idx += num_threads) { - out_vec(idx) = static_cast(out_smem(idx)); + if (local_v < kTileV && row_lane == 0 && v_global < kHeadDimV) { + float out_row = 0.f; +#pragma unroll + for (int lane = 0; lane < kThreadsPerVRow; ++lane) { + out_row += out_smem(local_v * kThreadsPerVRow + lane); + } + out_vec(v_global) = static_cast(out_row); } } }; From 68e8fa0e9c26a358bbf2d103673ba1eaae14d7ec Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Fri, 12 Jun 2026 12:50:55 +0000 Subject: [PATCH 13/35] Revert "Optimize qwen35 decode mainloop tiling" This reverts commit b38e9a2df1f69ad9fd9ca8512026a5e737d2f7ac. --- csrc/qwen35/decode/qwen35_decode_common.cuh | 4 +- .../decode/qwen35_scalar_kda_kernel.hpp | 21 ++-- .../decode/qwen35_scalar_kda_mainloop.hpp | 97 +++++++------------ 3 files changed, 42 insertions(+), 80 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_decode_common.cuh b/csrc/qwen35/decode/qwen35_decode_common.cuh index 2f78d22e..51ce1872 100644 --- a/csrc/qwen35/decode/qwen35_decode_common.cuh +++ b/csrc/qwen35/decode/qwen35_decode_common.cuh @@ -67,8 +67,8 @@ struct Qwen35DecodeLocalShape { static constexpr int kLayoutVec = 4; static constexpr int kLayoutThreads = kHeadDimQK / kLayoutVec; static constexpr int kKdaThreads = 128; - static constexpr int kKdaTileV = 32; - static constexpr int kKdaTileK = kHeadDimQK; + static constexpr int kKdaTileV = 16; + static constexpr int kKdaTileK = 16; static_assert(kLocalVHeads % kLocalQKHeads == 0); static_assert(kHeadDimQK == kHeadDimV); diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index ed67ccc6..9c91e8b3 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -57,7 +57,6 @@ struct Qwen35ScalarKdaDecodeKernel { alignas(16) float k_smem[kHeadDimQK]; alignas(16) scalar_t v_smem[kHeadDimV]; alignas(16) float norm_smem[2]; - alignas(16) float state_smem[kHeadDimQK * kTileV]; alignas(16) float proj_smem[kHeadDimV]; alignas(16) float out_smem[kHeadDimV]; }; @@ -67,12 +66,8 @@ struct Qwen35ScalarKdaDecodeKernel { } static dim3 grid_shape(int token_count) { - // One block owns one V tile for one (token_idx, hv) pair. - constexpr int kNumVTiles = (kHeadDimV + kTileV - 1) / kTileV; - return dim3( - static_cast(kNumVTiles), - static_cast(Shape::kLocalVHeads), - static_cast(token_count)); + // One block owns one (token_idx, hv) pair in the first implementation. + return dim3(static_cast(Shape::kLocalVHeads), static_cast(token_count), 1); } template @@ -89,9 +84,8 @@ struct Qwen35ScalarKdaDecodeKernel { scalar_t* __restrict__ out, int token_count, SharedStorage& storage) { - const int v_tile = static_cast(blockIdx.x); - const int hv = static_cast(blockIdx.y); - const int token_idx = static_cast(blockIdx.z); + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); const int tid = static_cast(threadIdx.x); if (token_idx >= token_count || hv >= kLocalVHeads) { return; @@ -189,7 +183,6 @@ struct Qwen35ScalarKdaDecodeKernel { state_vk, out_vec, storage, - v_tile * kTileV, tid, kThreads); } @@ -211,9 +204,8 @@ struct Qwen35ScalarKdaDecodeKernel { constexpr int kLocalKDim = Shape::kLocalKDim; constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; - const int v_tile = static_cast(blockIdx.x); - const int hv = static_cast(blockIdx.y); - const int token_idx = static_cast(blockIdx.z); + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); const int tid = static_cast(threadIdx.x); if (token_idx >= token_count || hv >= kLocalVHeads) { return; @@ -283,7 +275,6 @@ struct Qwen35ScalarKdaDecodeKernel { state_vk, out_vec, storage, - v_tile * kTileV, tid, kThreads); } diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp index 206cbf91..0c9401bc 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp @@ -56,8 +56,8 @@ struct Qwen35ScalarKdaDecodeMainloop { // 2. optimize proj / update / out reductions // 3. evaluate warp-specialized load/compute roles only after the fp32 path // is stable and measured - static constexpr int kTileV = 32; - static constexpr int kTileK = kHeadDimQK; + static constexpr int kTileV = 16; + static constexpr int kTileK = 16; static constexpr int kTilesPerV = kHeadDimV / kTileV; static constexpr int kTilesPerK = kHeadDimQK / kTileK; static constexpr int kWarpSize = 32; @@ -65,14 +65,10 @@ struct Qwen35ScalarKdaDecodeMainloop { static constexpr int kWarpsPerCta = 4; static constexpr int kRowsPerWarp = kWarpSize; static constexpr int kRowsPerThread = 1; - static constexpr int kThreadsPerVRow = 4; - static constexpr int kKPerThread = kHeadDimQK / kThreadsPerVRow; static_assert(kHeadDimV == 128); static_assert(kHeadDimQK == 128); static_assert(kWarpsPerCta * kRowsPerWarp == kHeadDimV); - static_assert(kTileV * kThreadsPerVRow == kWarpsPerCta * kWarpSize); - static_assert(kHeadDimQK % kThreadsPerVRow == 0); // First concrete decode threading plan: // @@ -362,7 +358,6 @@ struct Qwen35ScalarKdaDecodeMainloop { TensorHvk& state_vk, TensorOut& out_vec, SharedStorage& storage, - int v_tile_base, int tid, int num_threads) { // Decode organization: @@ -398,9 +393,6 @@ struct Qwen35ScalarKdaDecodeMainloop { auto k_smem = make_tensor(make_smem_ptr(storage.k_smem), make_layout(make_shape(Int{}))); auto v_smem = make_tensor(make_smem_ptr(storage.v_smem), make_layout(make_shape(Int{}))); auto norm_smem = make_tensor(make_smem_ptr(storage.norm_smem), make_layout(make_shape(Int<2>{}))); - auto state_smem = make_tensor( - make_smem_ptr(storage.state_smem), - make_layout(make_shape(Int{}, Int{}), make_stride(Int{}, Int<1>{}))); auto proj_smem = make_tensor(make_smem_ptr(storage.proj_smem), make_layout(make_shape(Int{}))); auto out_smem = make_tensor(make_smem_ptr(storage.out_smem), make_layout(make_shape(Int{}))); @@ -409,22 +401,11 @@ struct Qwen35ScalarKdaDecodeMainloop { q_smem(idx) = static_cast(q_vec(idx)); k_smem(idx) = static_cast(k_vec(idx)); } - for (int local_v = tid; local_v < kTileV; local_v += num_threads) { - const int v_global = v_tile_base + local_v; - if (v_global < kHeadDimV) { - v_smem(local_v) = v_vec(v_global); - } - } for (int idx = tid; idx < kHeadDimV; idx += num_threads) { + v_smem(idx) = v_vec(idx); proj_smem(idx) = 0.f; out_smem(idx) = 0.f; } - for (int idx = tid; idx < kHeadDimQK * kTileV; idx += num_threads) { - const int k_idx = idx / kTileV; - const int local_v = idx - k_idx * kTileV; - const int v_global = v_tile_base + local_v; - state_smem(k_idx, local_v) = v_global < kHeadDimV ? static_cast(state_vk(v_global, k_idx)) : 0.f; - } __syncthreads(); if (tid == 0) { @@ -448,55 +429,45 @@ struct Qwen35ScalarKdaDecodeMainloop { } __syncthreads(); - const int local_v = tid / kThreadsPerVRow; - const int row_lane = tid - local_v * kThreadsPerVRow; - const int v_global = v_tile_base + local_v; - if (local_v < kTileV && v_global < kHeadDimV) { - const int k_begin = row_lane * kKPerThread; - float proj_part = 0.f; -#pragma unroll - for (int kk = 0; kk < kKPerThread; ++kk) { - const int k_idx = k_begin + kk; - proj_part += state_smem(k_idx, local_v) * k_smem(k_idx); - } - proj_smem(tid) = proj_part; - } - __syncthreads(); + ThreadRowPlan row_plan = make_thread_row_plan(tid); - if (local_v < kTileV && row_lane == 0 && v_global < kHeadDimV) { + // First concrete ownership model: + // - each thread owns one full state row across all 128 K columns + // - the row is streamed tile-by-tile through registers + // - no cross-thread reduction is needed for proj/out because the full row + // stays with one thread for the duration of the token update + if (row_plan.owns_row) { float proj_row = 0.f; -#pragma unroll - for (int lane = 0; lane < kThreadsPerVRow; ++lane) { - proj_row += proj_smem(local_v * kThreadsPerVRow + lane); + for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { + TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; + RowTileProjPlan proj_plan = make_row_tile_proj_plan(coords, row_plan.warp_id, row_plan.lane_id); + proj_plan.v_row = row_plan.v_row; + proj_plan.owns_row = true; + proj_row += project_row_tile(state_vk, k_smem, proj_plan); } - const float v_val = static_cast(v_smem(local_v)); - proj_smem(local_v) = beta * (v_val - decay * proj_row); - } - __syncthreads(); - if (local_v < kTileV && v_global < kHeadDimV) { - const int k_begin = row_lane * kKPerThread; - const float v_new_row = proj_smem(local_v); - float out_part = 0.f; -#pragma unroll - for (int kk = 0; kk < kKPerThread; ++kk) { - const int k_idx = k_begin + kk; - const float state_new = decay * state_smem(k_idx, local_v) + v_new_row * k_smem(k_idx); - state_smem(k_idx, local_v) = state_new; - out_part += state_new * q_smem(k_idx); - state_vk(v_global, k_idx) = state_new; + proj_smem(row_plan.v_row) = proj_row; + + const float v_val = static_cast(v_smem(row_plan.v_row)); + const float decayed_proj_row = decay * proj_row; + const float v_new_row = beta * (v_val - decayed_proj_row); + + float out_row = 0.f; + for (int tile_k = 0; tile_k < kTilesPerK; ++tile_k) { + TileCoords coords = TileCoords{(row_plan.v_row / kTileV) * kTileV, tile_k * kTileK}; + RowTileUpdatePlan update_plan = make_row_tile_update_plan(coords, row_plan.warp_id, row_plan.lane_id); + update_plan.v_row = row_plan.v_row; + update_plan.owns_row = true; + out_row += update_and_output_row_tile( + state_vk, k_smem, q_smem, update_plan, decay, v_new_row); } - out_smem(local_v * kThreadsPerVRow + row_lane) = out_part; + + out_smem(row_plan.v_row) = out_row; } __syncthreads(); - if (local_v < kTileV && row_lane == 0 && v_global < kHeadDimV) { - float out_row = 0.f; -#pragma unroll - for (int lane = 0; lane < kThreadsPerVRow; ++lane) { - out_row += out_smem(local_v * kThreadsPerVRow + lane); - } - out_vec(v_global) = static_cast(out_row); + for (int idx = tid; idx < kHeadDimV; idx += num_threads) { + out_vec(idx) = static_cast(out_smem(idx)); } } }; From a9c2fa8dfbad63639551cc72640627f8fa22ac16 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sat, 13 Jun 2026 18:36:35 +0000 Subject: [PATCH 14/35] Optimize qwen35 long decode kernel --- benchmarks/profile_qwen35_decode.py | 180 ++++++++++++ .../decode/qwen35_scalar_kda_kernel.hpp | 269 ++++++++++++++++++ .../decode/qwen35_scalar_kda_mainloop.hpp | 23 +- 3 files changed, 469 insertions(+), 3 deletions(-) create mode 100644 benchmarks/profile_qwen35_decode.py diff --git a/benchmarks/profile_qwen35_decode.py b/benchmarks/profile_qwen35_decode.py new file mode 100644 index 00000000..b1993f5f --- /dev/null +++ b/benchmarks/profile_qwen35_decode.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Small Nsight Compute target for Qwen3.5 decode kernels. + +Use with `ncu --profile-from-start off` so only the decode loop bracketed by +cudaProfilerStart/Stop is collected. +""" + +from __future__ import annotations + +import argparse +import pathlib +import sys + +import torch + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +import cula.cudac as cula_cuda +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as GLOBAL_CONFIG +from cula.qwen35.common import Qwen35LinearAttentionConfig + + +def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: + if tp_size not in (1, 2, 4, 8): + raise ValueError(f"tp_size must be one of 1, 2, 4, 8, got {tp_size}") + return Qwen35LinearAttentionConfig( + hidden_size=GLOBAL_CONFIG.hidden_size // tp_size, + conv_kernel_size=GLOBAL_CONFIG.conv_kernel_size, + num_k_heads=GLOBAL_CONFIG.num_k_heads // tp_size, + num_v_heads=GLOBAL_CONFIG.num_v_heads // tp_size, + head_k_dim=GLOBAL_CONFIG.head_k_dim, + head_v_dim=GLOBAL_CONFIG.head_v_dim, + qkv_dtype=GLOBAL_CONFIG.qkv_dtype, + state_dtype=GLOBAL_CONFIG.state_dtype, + ) + + +def make_fused_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): + torch.manual_seed(seed) + mixed_qkv_conv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state = torch.randn( + tokens, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + state_work = torch.empty_like(state) + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + out = torch.empty(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) + return mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out + + +def make_native_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): + torch.manual_seed(seed) + q = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + k = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) + v = torch.randn(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) + a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) + dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + state = torch.randn( + tokens, + config.num_v_heads, + config.head_k_dim, + config.head_v_dim, + device=device, + dtype=config.state_dtype, + ) * 0.01 + state_work = torch.empty_like(state) + state_indices = torch.arange(tokens, device=device, dtype=torch.int32) + out = torch.empty_like(v) + return q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out + + +def profiler_start() -> None: + torch.cuda.profiler.start() + + +def profiler_stop() -> None: + torch.cuda.profiler.stop() + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--op", choices=("fused", "native"), default="fused") + parser.add_argument("--tokens", type=int, default=128) + parser.add_argument("--tp-size", type=int, default=1) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rep", type=int, default=3) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--device", type=int, default=0) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("No CUDA device is available") + + torch.cuda.set_device(args.device) + device = torch.device("cuda", torch.cuda.current_device()) + config = local_config_from_tp_size(args.tp_size) + + if args.op == "fused": + mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_fused_inputs( + args.tokens, device, args.seed, config + ) + + def run() -> None: + cula_cuda.qwen35_layout_scalar_kda_decode( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + state_work, + state_indices, + out, + ) + + else: + q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_native_inputs( + args.tokens, device, args.seed, config + ) + + def run() -> None: + cula_cuda.qwen35_scalar_kda_decode( + q, + k, + v, + a, + b, + A_log, + dt_bias, + state_work, + state_indices, + out, + ) + + for _ in range(args.warmup): + state_work.copy_(state) + run() + torch.cuda.synchronize() + + state_work.copy_(state) + torch.cuda.synchronize() + + print( + f"profile op={args.op} tokens={args.tokens} tp={args.tp_size} " + f"warmup={args.warmup} rep={args.rep} device={torch.cuda.get_device_name(device)}" + ) + profiler_start() + for _ in range(args.rep): + run() + torch.cuda.synchronize() + profiler_stop() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 9c91e8b3..31731bc7 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -23,6 +23,21 @@ namespace cula::qwen35::decode::kernel { using namespace cute; +template +CUTE_DEVICE void cp_async_ca_shared_global(void* smem_ptr, const void* gmem_ptr) { + static_assert(kBytes == 16, "Only 16-byte cp.async copies are supported here."); + const unsigned smem_addr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" ::"r"(smem_addr), "l"(gmem_ptr)); +} + +CUTE_DEVICE void cp_async_commit_group() { + asm volatile("cp.async.commit_group;\n" ::); +} + +CUTE_DEVICE void cp_async_wait_all() { + asm volatile("cp.async.wait_group 0;\n" ::); +} + template struct Qwen35ScalarKdaDecodeKernel { using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; @@ -323,6 +338,28 @@ void launch_qwen35_scalar_kda_decode_kernel( const int32_t* pool_idx, scalar_t* out, int token_count) { + if (token_count >= 64) { + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); + auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); + qwen35_scalar_kda_decode_kernel< + scalar_t, + kLocalQKHeads, + kLocalVHeads, + Qwen35ScalarKdaDecodeLongMainloop><<>>( + q_rep, + k_rep, + v, + a_kernel, + b_kernel, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); + return; + } + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); qwen35_scalar_kda_decode_kernel<<>>( @@ -364,6 +401,223 @@ __global__ void qwen35_layout_scalar_kda_decode_kernel( storage); } +template +__global__ void qwen35_layout_scalar_kda_decode_long_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; + constexpr int kWarpTileV = 32; + constexpr int kWarpSize = 32; + constexpr int kThreads = 128; + constexpr int kWarps = kThreads / kWarpSize; + constexpr int kPipeTileK = 16; + constexpr int kPipeStages = 2; + constexpr int kVecFloats = 4; + + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kHeadDimV % kWarpTileV == 0); + static_assert(kWarps * kWarpTileV == kHeadDimV); + static_assert(kHeadDimQK % kPipeTileK == 0); + + __shared__ float q_smem[kHeadDimQK]; + __shared__ float k_smem[kHeadDimQK]; + __shared__ float norm_smem[2 * kWarps]; + __shared__ float state_pipe[kPipeStages][kPipeTileK][kHeadDimV]; + + const int hv = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int tid = static_cast(threadIdx.x); + const int lane = tid & (kWarpSize - 1); + const int warp_id = tid / kWarpSize; + const int mapped_h = hv / kRepeatFactor; + const int v_row = warp_id * kWarpTileV + lane; + + auto qk_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimQK, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimV, Int<1>{})); + auto out_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + const scalar_t* q_src = mixed_qkv_conv; + const scalar_t* k_src = mixed_qkv_conv + kLocalQDim; + const scalar_t* v_src = mixed_qkv_conv + kLocalQDim + kLocalKDim; + + auto gQ = make_tensor(make_gmem_ptr(q_src), qk_src_layout); + auto gK = make_tensor(make_gmem_ptr(k_src), qk_src_layout); + auto gV = make_tensor(make_gmem_ptr(v_src), v_src_layout); + auto gO = make_tensor(make_gmem_ptr(out), out_layout); + auto gA = make_tensor(make_gmem_ptr(a), head_layout); + auto gB = make_tensor(make_gmem_ptr(b), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, mapped_h, _); + auto k_vec = gK(token_idx, mapped_h, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto state_vk = gH_vk(state_row, hv, _, _); + + const float a_val = static_cast(gA(token_idx, hv)); + const float b_val = static_cast(gB(token_idx, hv)); + const float g = -expf(static_cast(gAlog(hv))) * + Qwen35ScalarKdaDecodeMainloop::softplusf_approx(a_val + static_cast(gDt(hv))); + const float decay = expf(g); + const float beta = 1.f / (1.f + expf(-b_val)); + + const float q_raw = static_cast(q_vec(tid)); + const float k_raw = static_cast(k_vec(tid)); + q_smem[tid] = q_raw; + k_smem[tid] = k_raw; + + float q_norm_sq = q_raw * q_raw; + float k_norm_sq = k_raw * k_raw; + q_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_norm_sq); + k_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_norm_sq); + if (lane == 0) { + norm_smem[warp_id] = q_norm_sq; + norm_smem[kWarps + warp_id] = k_norm_sq; + } + __syncthreads(); + + float q_block_sum = lane < kWarps ? norm_smem[lane] : 0.f; + float k_block_sum = lane < kWarps ? norm_smem[kWarps + lane] : 0.f; + if (warp_id == 0) { + q_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_block_sum); + k_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_block_sum); + if (lane == 0) { + norm_smem[0] = rsqrtf(q_block_sum + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); + norm_smem[1] = rsqrtf(k_block_sum + 1e-6f); + } + } + __syncthreads(); + + q_smem[tid] = q_smem[tid] * norm_smem[0]; + k_smem[tid] = k_smem[tid] * norm_smem[1]; + __syncthreads(); + + auto load_state_pipe_tile = [&](int stage, int k_base) { +#pragma unroll 1 + for (int elem = tid * kVecFloats; elem < kPipeTileK * kHeadDimV; elem += kThreads * kVecFloats) { + const int k_local = elem / kHeadDimV; + const int v_base = elem - k_local * kHeadDimV; + cp_async_ca_shared_global<16>( + &state_pipe[stage][k_local][v_base], + &state_vk(v_base, k_base + k_local)); + } + }; + + float proj_row = 0.f; + int pipe_stage = 0; + load_state_pipe_tile(pipe_stage, 0); + cp_async_commit_group(); + +#pragma unroll 1 + for (int k_base = 0; k_base < kHeadDimQK; k_base += kPipeTileK) { + cp_async_wait_all(); + __syncthreads(); + + const int next_k_base = k_base + kPipeTileK; + const int next_stage = pipe_stage ^ 1; + if (next_k_base < kHeadDimQK) { + load_state_pipe_tile(next_stage, next_k_base); + cp_async_commit_group(); + } + +#pragma unroll + for (int k_local = 0; k_local < kPipeTileK; k_local += 4) { + const float state0 = state_pipe[pipe_stage][k_local + 0][v_row]; + const float state1 = state_pipe[pipe_stage][k_local + 1][v_row]; + const float state2 = state_pipe[pipe_stage][k_local + 2][v_row]; + const float state3 = state_pipe[pipe_stage][k_local + 3][v_row]; + proj_row += state0 * k_smem[k_base + k_local + 0]; + proj_row += state1 * k_smem[k_base + k_local + 1]; + proj_row += state2 * k_smem[k_base + k_local + 2]; + proj_row += state3 * k_smem[k_base + k_local + 3]; + } + __syncthreads(); + pipe_stage = next_stage; + } + + const float v_val = static_cast(v_vec(v_row)); + const float v_new_row = beta * (v_val - decay * proj_row); + + float out_row = 0.f; +#pragma unroll 1 + for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 4) { + const float state_new0 = decay * static_cast(state_vk(v_row, k_idx + 0)) + v_new_row * k_smem[k_idx + 0]; + const float state_new1 = decay * static_cast(state_vk(v_row, k_idx + 1)) + v_new_row * k_smem[k_idx + 1]; + const float state_new2 = decay * static_cast(state_vk(v_row, k_idx + 2)) + v_new_row * k_smem[k_idx + 2]; + const float state_new3 = decay * static_cast(state_vk(v_row, k_idx + 3)) + v_new_row * k_smem[k_idx + 3]; + state_vk(v_row, k_idx + 0) = state_new0; + state_vk(v_row, k_idx + 1) = state_new1; + state_vk(v_row, k_idx + 2) = state_new2; + state_vk(v_row, k_idx + 3) = state_new3; + out_row += state_new0 * q_smem[k_idx + 0]; + out_row += state_new1 * q_smem[k_idx + 1]; + out_row += state_new2 * q_smem[k_idx + 2]; + out_row += state_new3 * q_smem[k_idx + 3]; + } + out_vec(v_row) = static_cast(out_row); +} + +template +void launch_qwen35_layout_scalar_kda_decode_long_kernel( + cudaStream_t stream, + const scalar_t* mixed_qkv_conv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + float* recurrent_state, + const int32_t* pool_idx, + scalar_t* out, + int token_count) { + constexpr int kWarpTileV = 32; + (void)kWarpTileV; + dim3 grid(kLocalVHeads, token_count, 1); + dim3 block(128, 1, 1); + qwen35_layout_scalar_kda_decode_long_kernel<<>>( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); +} + template > void launch_qwen35_layout_scalar_kda_decode_kernel( cudaStream_t stream, @@ -376,6 +630,21 @@ void launch_qwen35_layout_scalar_kda_decode_kernel( const int32_t* pool_idx, scalar_t* out, int token_count) { + if (token_count >= 64) { + launch_qwen35_layout_scalar_kda_decode_long_kernel( + stream, + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); + return; + } + auto grid = Qwen35ScalarKdaDecodeKernel::grid_shape(token_count); auto block = Qwen35ScalarKdaDecodeKernel::block_shape(); qwen35_layout_scalar_kda_decode_kernel<<>>( diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp index 0c9401bc..74460077 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_mainloop.hpp @@ -25,7 +25,7 @@ namespace cula::qwen35::decode::kernel { using namespace cute; -template +template struct Qwen35ScalarKdaDecodeMainloop { // Decode design decision: // - recurrent_state remains fp32 both physically and mathematically @@ -56,8 +56,8 @@ struct Qwen35ScalarKdaDecodeMainloop { // 2. optimize proj / update / out reductions // 3. evaluate warp-specialized load/compute roles only after the fp32 path // is stable and measured - static constexpr int kTileV = 16; - static constexpr int kTileK = 16; + static constexpr int kTileV = kTileV_; + static constexpr int kTileK = kTileK_; static constexpr int kTilesPerV = kHeadDimV / kTileV; static constexpr int kTilesPerK = kHeadDimQK / kTileK; static constexpr int kWarpSize = 32; @@ -68,6 +68,8 @@ struct Qwen35ScalarKdaDecodeMainloop { static_assert(kHeadDimV == 128); static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV % kTileV == 0); + static_assert(kHeadDimQK % kTileK == 0); static_assert(kWarpsPerCta * kRowsPerWarp == kHeadDimV); // First concrete decode threading plan: @@ -336,6 +338,15 @@ struct Qwen35ScalarKdaDecodeMainloop { return x > 20.f ? x : log1pf(expf(x)); } + CUTE_DEVICE static float warp_sum(float value) { + constexpr unsigned int kFullMask = 0xffffffffu; +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(kFullMask, value, offset); + } + return value; + } + template < typename TensorQ, typename TensorK, @@ -472,4 +483,10 @@ struct Qwen35ScalarKdaDecodeMainloop { } }; +template +struct Qwen35ScalarKdaDecodeLongMainloop : public Qwen35ScalarKdaDecodeMainloop { + using Base = Qwen35ScalarKdaDecodeMainloop; + using Base::run; +}; + } // namespace cula::qwen35::decode::kernel From 7d60b3b9016d4654ba45073e747eaee4770b76d4 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sat, 13 Jun 2026 19:04:08 +0000 Subject: [PATCH 15/35] Pipeline qwen35 long decode update --- .../decode/qwen35_scalar_kda_kernel.hpp | 67 ++++++++++++++----- 1 file changed, 51 insertions(+), 16 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 31731bc7..1152e986 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -424,6 +424,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( constexpr int kPipeTileK = 16; constexpr int kPipeStages = 2; constexpr int kVecFloats = 4; + constexpr int kStatePipeStrideV = kHeadDimV + 4; static_assert(kLocalQKHeads == Shape::kLocalQKHeads); static_assert(kHeadDimQK == 128); @@ -435,7 +436,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( __shared__ float q_smem[kHeadDimQK]; __shared__ float k_smem[kHeadDimQK]; __shared__ float norm_smem[2 * kWarps]; - __shared__ float state_pipe[kPipeStages][kPipeTileK][kHeadDimV]; + __shared__ float state_pipe[kPipeStages][kPipeTileK][kStatePipeStrideV]; const int hv = static_cast(blockIdx.x); const int token_idx = static_cast(blockIdx.y); @@ -525,6 +526,21 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( k_smem[tid] = k_smem[tid] * norm_smem[1]; __syncthreads(); + float qk_dot = q_smem[tid] * k_smem[tid]; + qk_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_dot); + if (lane == 0) { + norm_smem[warp_id] = qk_dot; + } + __syncthreads(); + float qk_block_sum = lane < kWarps ? norm_smem[lane] : 0.f; + if (warp_id == 0) { + qk_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_block_sum); + if (lane == 0) { + norm_smem[2] = qk_block_sum; + } + } + __syncthreads(); + auto load_state_pipe_tile = [&](int stage, int k_base) { #pragma unroll 1 for (int elem = tid * kVecFloats; elem < kPipeTileK * kHeadDimV; elem += kThreads * kVecFloats) { @@ -537,6 +553,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( }; float proj_row = 0.f; + float out_old_row = 0.f; int pipe_stage = 0; load_state_pipe_tile(pipe_stage, 0); cp_async_commit_group(); @@ -563,6 +580,10 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( proj_row += state1 * k_smem[k_base + k_local + 1]; proj_row += state2 * k_smem[k_base + k_local + 2]; proj_row += state3 * k_smem[k_base + k_local + 3]; + out_old_row += state0 * q_smem[k_base + k_local + 0]; + out_old_row += state1 * q_smem[k_base + k_local + 1]; + out_old_row += state2 * q_smem[k_base + k_local + 2]; + out_old_row += state3 * q_smem[k_base + k_local + 3]; } __syncthreads(); pipe_stage = next_stage; @@ -570,24 +591,38 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( const float v_val = static_cast(v_vec(v_row)); const float v_new_row = beta * (v_val - decay * proj_row); + out_vec(v_row) = static_cast(decay * out_old_row + v_new_row * norm_smem[2]); + + pipe_stage = 0; + load_state_pipe_tile(pipe_stage, 0); + cp_async_commit_group(); - float out_row = 0.f; #pragma unroll 1 - for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 4) { - const float state_new0 = decay * static_cast(state_vk(v_row, k_idx + 0)) + v_new_row * k_smem[k_idx + 0]; - const float state_new1 = decay * static_cast(state_vk(v_row, k_idx + 1)) + v_new_row * k_smem[k_idx + 1]; - const float state_new2 = decay * static_cast(state_vk(v_row, k_idx + 2)) + v_new_row * k_smem[k_idx + 2]; - const float state_new3 = decay * static_cast(state_vk(v_row, k_idx + 3)) + v_new_row * k_smem[k_idx + 3]; - state_vk(v_row, k_idx + 0) = state_new0; - state_vk(v_row, k_idx + 1) = state_new1; - state_vk(v_row, k_idx + 2) = state_new2; - state_vk(v_row, k_idx + 3) = state_new3; - out_row += state_new0 * q_smem[k_idx + 0]; - out_row += state_new1 * q_smem[k_idx + 1]; - out_row += state_new2 * q_smem[k_idx + 2]; - out_row += state_new3 * q_smem[k_idx + 3]; + for (int k_base = 0; k_base < kHeadDimQK; k_base += kPipeTileK) { + cp_async_wait_all(); + __syncthreads(); + + const int next_k_base = k_base + kPipeTileK; + const int next_stage = pipe_stage ^ 1; + if (next_k_base < kHeadDimQK) { + load_state_pipe_tile(next_stage, next_k_base); + cp_async_commit_group(); + } + +#pragma unroll + for (int k_local = 0; k_local < kPipeTileK; k_local += 4) { + const float state_new0 = decay * state_pipe[pipe_stage][k_local + 0][v_row] + v_new_row * k_smem[k_base + k_local + 0]; + const float state_new1 = decay * state_pipe[pipe_stage][k_local + 1][v_row] + v_new_row * k_smem[k_base + k_local + 1]; + const float state_new2 = decay * state_pipe[pipe_stage][k_local + 2][v_row] + v_new_row * k_smem[k_base + k_local + 2]; + const float state_new3 = decay * state_pipe[pipe_stage][k_local + 3][v_row] + v_new_row * k_smem[k_base + k_local + 3]; + state_vk(v_row, k_base + k_local + 0) = state_new0; + state_vk(v_row, k_base + k_local + 1) = state_new1; + state_vk(v_row, k_base + k_local + 2) = state_new2; + state_vk(v_row, k_base + k_local + 3) = state_new3; + } + __syncthreads(); + pipe_stage = next_stage; } - out_vec(v_row) = static_cast(out_row); } template From e99eda5044beadba0eaec64ea009405f5d55d407 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 15 Jun 2026 02:54:30 +0000 Subject: [PATCH 16/35] Optimize qwen35 long decode ILP --- .../decode/qwen35_scalar_kda_kernel.hpp | 128 +++++++++++++----- 1 file changed, 93 insertions(+), 35 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 1152e986..6fbf30b3 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -38,6 +38,10 @@ CUTE_DEVICE void cp_async_wait_all() { asm volatile("cp.async.wait_group 0;\n" ::); } +CUTE_DEVICE void cp_async_wait_group_1() { + asm volatile("cp.async.wait_group 1;\n" ::); +} + template struct Qwen35ScalarKdaDecodeKernel { using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; @@ -401,7 +405,7 @@ __global__ void qwen35_layout_scalar_kda_decode_kernel( storage); } -template +template __global__ void qwen35_layout_scalar_kda_decode_long_kernel( const scalar_t* __restrict__ mixed_qkv_conv, const scalar_t* __restrict__ a, @@ -421,7 +425,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( constexpr int kWarpSize = 32; constexpr int kThreads = 128; constexpr int kWarps = kThreads / kWarpSize; - constexpr int kPipeTileK = 16; + constexpr int kPipeTileK = kPipeTileK_; constexpr int kPipeStages = 2; constexpr int kVecFloats = 4; constexpr int kStatePipeStrideV = kHeadDimV + 4; @@ -432,6 +436,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( static_assert(kHeadDimV % kWarpTileV == 0); static_assert(kWarps * kWarpTileV == kHeadDimV); static_assert(kHeadDimQK % kPipeTileK == 0); + static_assert(kPipeTileK == 16 || kPipeTileK == 32); __shared__ float q_smem[kHeadDimQK]; __shared__ float k_smem[kHeadDimQK]; @@ -522,11 +527,12 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( } __syncthreads(); - q_smem[tid] = q_smem[tid] * norm_smem[0]; - k_smem[tid] = k_smem[tid] * norm_smem[1]; - __syncthreads(); + const float q_normed = q_raw * norm_smem[0]; + const float k_normed = k_raw * norm_smem[1]; + q_smem[tid] = q_normed; + k_smem[tid] = k_normed; - float qk_dot = q_smem[tid] * k_smem[tid]; + float qk_dot = q_normed * k_normed; qk_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_dot); if (lane == 0) { norm_smem[warp_id] = qk_dot; @@ -552,22 +558,40 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( } }; - float proj_row = 0.f; - float out_old_row = 0.f; + float proj_acc0 = 0.f; + float proj_acc1 = 0.f; + float proj_acc2 = 0.f; + float proj_acc3 = 0.f; + float out_acc0 = 0.f; + float out_acc1 = 0.f; + float out_acc2 = 0.f; + float out_acc3 = 0.f; int pipe_stage = 0; - load_state_pipe_tile(pipe_stage, 0); + load_state_pipe_tile(0, 0); cp_async_commit_group(); + if (kPipeTileK < kHeadDimQK) { + load_state_pipe_tile(1, kPipeTileK); + cp_async_commit_group(); + } #pragma unroll 1 for (int k_base = 0; k_base < kHeadDimQK; k_base += kPipeTileK) { - cp_async_wait_all(); - __syncthreads(); - const int next_k_base = k_base + kPipeTileK; const int next_stage = pipe_stage ^ 1; + const int prefetch_k_base = k_base + 2 * kPipeTileK; if (next_k_base < kHeadDimQK) { - load_state_pipe_tile(next_stage, next_k_base); - cp_async_commit_group(); + cp_async_wait_group_1(); + } else { + cp_async_wait_all(); + } + __syncthreads(); + + float q_regs[kPipeTileK]; + float k_regs[kPipeTileK]; +#pragma unroll + for (int kk = 0; kk < kPipeTileK; ++kk) { + q_regs[kk] = q_smem[k_base + kk]; + k_regs[kk] = k_smem[k_base + kk]; } #pragma unroll @@ -576,51 +600,71 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( const float state1 = state_pipe[pipe_stage][k_local + 1][v_row]; const float state2 = state_pipe[pipe_stage][k_local + 2][v_row]; const float state3 = state_pipe[pipe_stage][k_local + 3][v_row]; - proj_row += state0 * k_smem[k_base + k_local + 0]; - proj_row += state1 * k_smem[k_base + k_local + 1]; - proj_row += state2 * k_smem[k_base + k_local + 2]; - proj_row += state3 * k_smem[k_base + k_local + 3]; - out_old_row += state0 * q_smem[k_base + k_local + 0]; - out_old_row += state1 * q_smem[k_base + k_local + 1]; - out_old_row += state2 * q_smem[k_base + k_local + 2]; - out_old_row += state3 * q_smem[k_base + k_local + 3]; + proj_acc0 += state0 * k_regs[k_local + 0]; + proj_acc1 += state1 * k_regs[k_local + 1]; + proj_acc2 += state2 * k_regs[k_local + 2]; + proj_acc3 += state3 * k_regs[k_local + 3]; + out_acc0 += state0 * q_regs[k_local + 0]; + out_acc1 += state1 * q_regs[k_local + 1]; + out_acc2 += state2 * q_regs[k_local + 2]; + out_acc3 += state3 * q_regs[k_local + 3]; + } + if (prefetch_k_base < kHeadDimQK) { + __syncthreads(); + load_state_pipe_tile(pipe_stage, prefetch_k_base); + cp_async_commit_group(); } - __syncthreads(); pipe_stage = next_stage; } + const float proj_row = (proj_acc0 + proj_acc1) + (proj_acc2 + proj_acc3); + const float out_old_row = (out_acc0 + out_acc1) + (out_acc2 + out_acc3); const float v_val = static_cast(v_vec(v_row)); const float v_new_row = beta * (v_val - decay * proj_row); out_vec(v_row) = static_cast(decay * out_old_row + v_new_row * norm_smem[2]); pipe_stage = 0; - load_state_pipe_tile(pipe_stage, 0); + load_state_pipe_tile(0, 0); cp_async_commit_group(); + if (kPipeTileK < kHeadDimQK) { + load_state_pipe_tile(1, kPipeTileK); + cp_async_commit_group(); + } #pragma unroll 1 for (int k_base = 0; k_base < kHeadDimQK; k_base += kPipeTileK) { - cp_async_wait_all(); - __syncthreads(); - const int next_k_base = k_base + kPipeTileK; const int next_stage = pipe_stage ^ 1; + const int prefetch_k_base = k_base + 2 * kPipeTileK; if (next_k_base < kHeadDimQK) { - load_state_pipe_tile(next_stage, next_k_base); - cp_async_commit_group(); + cp_async_wait_group_1(); + } else { + cp_async_wait_all(); + } + __syncthreads(); + + float k_regs[kPipeTileK]; +#pragma unroll + for (int kk = 0; kk < kPipeTileK; ++kk) { + k_regs[kk] = k_smem[k_base + kk]; } #pragma unroll for (int k_local = 0; k_local < kPipeTileK; k_local += 4) { - const float state_new0 = decay * state_pipe[pipe_stage][k_local + 0][v_row] + v_new_row * k_smem[k_base + k_local + 0]; - const float state_new1 = decay * state_pipe[pipe_stage][k_local + 1][v_row] + v_new_row * k_smem[k_base + k_local + 1]; - const float state_new2 = decay * state_pipe[pipe_stage][k_local + 2][v_row] + v_new_row * k_smem[k_base + k_local + 2]; - const float state_new3 = decay * state_pipe[pipe_stage][k_local + 3][v_row] + v_new_row * k_smem[k_base + k_local + 3]; + const float state_new0 = decay * state_pipe[pipe_stage][k_local + 0][v_row] + v_new_row * k_regs[k_local + 0]; + const float state_new1 = decay * state_pipe[pipe_stage][k_local + 1][v_row] + v_new_row * k_regs[k_local + 1]; + const float state_new2 = decay * state_pipe[pipe_stage][k_local + 2][v_row] + v_new_row * k_regs[k_local + 2]; + const float state_new3 = decay * state_pipe[pipe_stage][k_local + 3][v_row] + v_new_row * k_regs[k_local + 3]; state_vk(v_row, k_base + k_local + 0) = state_new0; state_vk(v_row, k_base + k_local + 1) = state_new1; state_vk(v_row, k_base + k_local + 2) = state_new2; state_vk(v_row, k_base + k_local + 3) = state_new3; } - __syncthreads(); + if (prefetch_k_base < kHeadDimQK) { + __syncthreads(); + load_state_pipe_tile(pipe_stage, prefetch_k_base); + cp_async_commit_group(); + } pipe_stage = next_stage; } } @@ -641,7 +685,21 @@ void launch_qwen35_layout_scalar_kda_decode_long_kernel( (void)kWarpTileV; dim3 grid(kLocalVHeads, token_count, 1); dim3 block(128, 1, 1); - qwen35_layout_scalar_kda_decode_long_kernel<<>>( + if (token_count == 64 || token_count == 128) { + qwen35_layout_scalar_kda_decode_long_kernel<<>>( + mixed_qkv_conv, + a, + b, + A_log, + dt_bias, + recurrent_state, + pool_idx, + out, + token_count); + return; + } + + qwen35_layout_scalar_kda_decode_long_kernel<<>>( mixed_qkv_conv, a, b, From a3c313bfe98c824074c2e7ba27908a470fcf4333 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sat, 11 Jul 2026 15:03:14 +0000 Subject: [PATCH 17/35] Optimize qwen35 long decode V32 staging --- .../decode/qwen35_scalar_kda_kernel.hpp | 248 +++++++++++++++++- 1 file changed, 245 insertions(+), 3 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index 6fbf30b3..a3758725 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -667,6 +667,246 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( } pipe_stage = next_stage; } + +} + + +template +__global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( + const scalar_t* __restrict__ mixed_qkv_conv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ recurrent_state, + const int32_t* __restrict__ pool_idx, + scalar_t* __restrict__ out, + int token_count) { + using Shape = cula::qwen35::decode::Qwen35DecodeLocalShape; + constexpr int kRepeatFactor = Shape::kRepeatFactor; + constexpr int kLocalQDim = Shape::kLocalQDim; + constexpr int kLocalKDim = Shape::kLocalKDim; + constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; + constexpr int kWarpSize = 32; + constexpr int kThreads = 32; + constexpr int kTileV = 32; + constexpr int kKPerThread = kHeadDimQK / kThreads; + constexpr int kVecFloats = 4; + + static_assert(kLocalQKHeads == Shape::kLocalQKHeads); + static_assert(kHeadDimQK == 128); + static_assert(kHeadDimV == 128); + static_assert(kHeadDimQK % kThreads == 0); + + __shared__ float q_smem[kHeadDimQK]; + __shared__ float k_smem[kHeadDimQK]; + __shared__ float state_smem[kHeadDimQK][kTileV]; + __shared__ float norm_smem[3]; + + const int hv_tile = static_cast(blockIdx.x); + const int token_idx = static_cast(blockIdx.y); + const int hv = hv_tile >> 2; + const int v_tile = hv_tile & 3; + const int tid = static_cast(threadIdx.x); + const int lane = tid & (kWarpSize - 1); + const int mapped_h = hv / kRepeatFactor; + const int v_row = v_tile * kTileV + lane; + + auto qk_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimQK, Int<1>{})); + auto v_src_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalMixedQKVDim, kHeadDimV, Int<1>{})); + auto out_layout = make_layout( + make_shape(token_count, Int{}, Int{}), + make_stride(kLocalVHeads * kHeadDimV, kHeadDimV, Int<1>{})); + auto head_layout = make_layout( + make_shape(token_count, Int{}), + make_stride(kLocalVHeads, Int<1>{})); + auto hv_layout = make_layout(make_shape(Int{}), make_stride(Int<1>{})); + auto state_layout_vk = make_layout( + make_shape(_, Int{}, Int{}, Int{}), + make_stride(Int{} * kHeadDimQK * kHeadDimV, kHeadDimQK * kHeadDimV, Int<1>{}, kHeadDimV)); + + const scalar_t* q_src = mixed_qkv_conv; + const scalar_t* k_src = mixed_qkv_conv + kLocalQDim; + const scalar_t* v_src = mixed_qkv_conv + kLocalQDim + kLocalKDim; + + auto gQ = make_tensor(make_gmem_ptr(q_src), qk_src_layout); + auto gK = make_tensor(make_gmem_ptr(k_src), qk_src_layout); + auto gV = make_tensor(make_gmem_ptr(v_src), v_src_layout); + auto gO = make_tensor(make_gmem_ptr(out), out_layout); + auto gA = make_tensor(make_gmem_ptr(a), head_layout); + auto gB = make_tensor(make_gmem_ptr(b), head_layout); + auto gAlog = make_tensor(make_gmem_ptr(A_log), hv_layout); + auto gDt = make_tensor(make_gmem_ptr(dt_bias), hv_layout); + auto gH_vk = make_tensor(make_gmem_ptr(recurrent_state), state_layout_vk); + + const int state_row = pool_idx[token_idx]; + if (state_row < 0) { + return; + } + + auto q_vec = gQ(token_idx, mapped_h, _); + auto k_vec = gK(token_idx, mapped_h, _); + auto v_vec = gV(token_idx, hv, _); + auto out_vec = gO(token_idx, hv, _); + auto state_vk = gH_vk(state_row, hv, _, _); + + float q_norm_sq = 0.f; + float k_norm_sq = 0.f; +#pragma unroll + for (int i = 0; i < kKPerThread; ++i) { + const int k_idx = i * kThreads + tid; + const float q_raw = static_cast(q_vec(k_idx)); + const float k_raw = static_cast(k_vec(k_idx)); + q_smem[k_idx] = q_raw; + k_smem[k_idx] = k_raw; + q_norm_sq += q_raw * q_raw; + k_norm_sq += k_raw * k_raw; + } + q_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_norm_sq); + k_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_norm_sq); + if (lane == 0) { + norm_smem[0] = rsqrtf(q_norm_sq + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); + norm_smem[1] = rsqrtf(k_norm_sq + 1e-6f); + } + __syncthreads(); + + float qk_dot = 0.f; +#pragma unroll + for (int i = 0; i < kKPerThread; ++i) { + const int k_idx = i * kThreads + tid; + const float q_normed = q_smem[k_idx] * norm_smem[0]; + const float k_normed = k_smem[k_idx] * norm_smem[1]; + q_smem[k_idx] = q_normed; + k_smem[k_idx] = k_normed; + qk_dot += q_normed * k_normed; + } + qk_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_dot); + if (lane == 0) { + norm_smem[2] = qk_dot; + } + __syncthreads(); + + const float a_val = static_cast(gA(token_idx, hv)); + const float b_val = static_cast(gB(token_idx, hv)); + const float g = -expf(static_cast(gAlog(hv))) * + Qwen35ScalarKdaDecodeMainloop::softplusf_approx(a_val + static_cast(gDt(hv))); + const float decay = expf(g); + const float beta = 1.f / (1.f + expf(-b_val)); + +#pragma unroll 1 + for (int elem = tid * kVecFloats; elem < kHeadDimQK * kTileV; elem += kThreads * kVecFloats) { + const int k_idx = elem / kTileV; + const int v_base = elem - k_idx * kTileV; + cp_async_ca_shared_global<16>( + &state_smem[k_idx][v_base], + &state_vk(v_tile * kTileV + v_base, k_idx)); + } + cp_async_commit_group(); + cp_async_wait_all(); + __syncthreads(); + + float proj_acc0 = 0.f; + float proj_acc1 = 0.f; + float proj_acc2 = 0.f; + float proj_acc3 = 0.f; + float proj_acc4 = 0.f; + float proj_acc5 = 0.f; + float proj_acc6 = 0.f; + float proj_acc7 = 0.f; + float out_acc0 = 0.f; + float out_acc1 = 0.f; + float out_acc2 = 0.f; + float out_acc3 = 0.f; + float out_acc4 = 0.f; + float out_acc5 = 0.f; + float out_acc6 = 0.f; + float out_acc7 = 0.f; +#pragma unroll 1 + for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 8) { + const float state0 = state_smem[k_idx + 0][lane]; + const float state1 = state_smem[k_idx + 1][lane]; + const float state2 = state_smem[k_idx + 2][lane]; + const float state3 = state_smem[k_idx + 3][lane]; + const float state4 = state_smem[k_idx + 4][lane]; + const float state5 = state_smem[k_idx + 5][lane]; + const float state6 = state_smem[k_idx + 6][lane]; + const float state7 = state_smem[k_idx + 7][lane]; + const float k0 = k_smem[k_idx + 0]; + const float k1 = k_smem[k_idx + 1]; + const float k2 = k_smem[k_idx + 2]; + const float k3 = k_smem[k_idx + 3]; + const float k4 = k_smem[k_idx + 4]; + const float k5 = k_smem[k_idx + 5]; + const float k6 = k_smem[k_idx + 6]; + const float k7 = k_smem[k_idx + 7]; + const float q0 = q_smem[k_idx + 0]; + const float q1 = q_smem[k_idx + 1]; + const float q2 = q_smem[k_idx + 2]; + const float q3 = q_smem[k_idx + 3]; + const float q4 = q_smem[k_idx + 4]; + const float q5 = q_smem[k_idx + 5]; + const float q6 = q_smem[k_idx + 6]; + const float q7 = q_smem[k_idx + 7]; + proj_acc0 += state0 * k0; + proj_acc1 += state1 * k1; + proj_acc2 += state2 * k2; + proj_acc3 += state3 * k3; + proj_acc4 += state4 * k4; + proj_acc5 += state5 * k5; + proj_acc6 += state6 * k6; + proj_acc7 += state7 * k7; + out_acc0 += state0 * q0; + out_acc1 += state1 * q1; + out_acc2 += state2 * q2; + out_acc3 += state3 * q3; + out_acc4 += state4 * q4; + out_acc5 += state5 * q5; + out_acc6 += state6 * q6; + out_acc7 += state7 * q7; + } + + const float proj_row = + ((proj_acc0 + proj_acc1) + (proj_acc2 + proj_acc3)) + + ((proj_acc4 + proj_acc5) + (proj_acc6 + proj_acc7)); + const float out_old_row = + ((out_acc0 + out_acc1) + (out_acc2 + out_acc3)) + + ((out_acc4 + out_acc5) + (out_acc6 + out_acc7)); + const float v_val = static_cast(v_vec(v_row)); + const float v_new_row = beta * (v_val - decay * proj_row); + out_vec(v_row) = static_cast(decay * out_old_row + v_new_row * norm_smem[2]); + +#pragma unroll 1 + for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 8) { + const float state0 = state_smem[k_idx + 0][lane]; + const float state1 = state_smem[k_idx + 1][lane]; + const float state2 = state_smem[k_idx + 2][lane]; + const float state3 = state_smem[k_idx + 3][lane]; + const float state4 = state_smem[k_idx + 4][lane]; + const float state5 = state_smem[k_idx + 5][lane]; + const float state6 = state_smem[k_idx + 6][lane]; + const float state7 = state_smem[k_idx + 7][lane]; + const float state_new0 = decay * state0 + v_new_row * k_smem[k_idx + 0]; + const float state_new1 = decay * state1 + v_new_row * k_smem[k_idx + 1]; + const float state_new2 = decay * state2 + v_new_row * k_smem[k_idx + 2]; + const float state_new3 = decay * state3 + v_new_row * k_smem[k_idx + 3]; + const float state_new4 = decay * state4 + v_new_row * k_smem[k_idx + 4]; + const float state_new5 = decay * state5 + v_new_row * k_smem[k_idx + 5]; + const float state_new6 = decay * state6 + v_new_row * k_smem[k_idx + 6]; + const float state_new7 = decay * state7 + v_new_row * k_smem[k_idx + 7]; + state_vk(v_row, k_idx + 0) = state_new0; + state_vk(v_row, k_idx + 1) = state_new1; + state_vk(v_row, k_idx + 2) = state_new2; + state_vk(v_row, k_idx + 3) = state_new3; + state_vk(v_row, k_idx + 4) = state_new4; + state_vk(v_row, k_idx + 5) = state_new5; + state_vk(v_row, k_idx + 6) = state_new6; + state_vk(v_row, k_idx + 7) = state_new7; + } + } template @@ -683,10 +923,10 @@ void launch_qwen35_layout_scalar_kda_decode_long_kernel( int token_count) { constexpr int kWarpTileV = 32; (void)kWarpTileV; - dim3 grid(kLocalVHeads, token_count, 1); - dim3 block(128, 1, 1); if (token_count == 64 || token_count == 128) { - qwen35_layout_scalar_kda_decode_long_kernel<<>>( + dim3 grid(kLocalVHeads * 4, token_count, 1); + dim3 block(32, 1, 1); + qwen35_layout_scalar_kda_decode_long_v32_kernel<<>>( mixed_qkv_conv, a, b, @@ -699,6 +939,8 @@ void launch_qwen35_layout_scalar_kda_decode_long_kernel( return; } + dim3 grid(kLocalVHeads, token_count, 1); + dim3 block(128, 1, 1); qwen35_layout_scalar_kda_decode_long_kernel<<>>( mixed_qkv_conv, a, From 3bd7c80b25b4fd92e968c41f46565daa8b3204e2 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sat, 11 Jul 2026 17:35:13 +0000 Subject: [PATCH 18/35] Optimize qwen35 long decode state staging --- .../decode/qwen35_scalar_kda_kernel.hpp | 130 +++++++++++++----- tests/test_qwen35_decode.py | 38 +++++ 2 files changed, 137 insertions(+), 31 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index a3758725..e84b4689 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -18,6 +18,7 @@ #include "qwen35_scalar_kda_mainloop.hpp" #include +#include namespace cula::qwen35::decode::kernel { @@ -30,6 +31,28 @@ CUTE_DEVICE void cp_async_ca_shared_global(void* smem_ptr, const void* gmem_ptr) asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" ::"r"(smem_addr), "l"(gmem_ptr)); } +CUTE_DEVICE void cp_async_bulk_shared_global( + void* smem_ptr, + const void* gmem_ptr, + uint32_t bytes, + cutlass::arch::ClusterTransactionBarrier::ValueType* barrier) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + const uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + const uint32_t barrier_addr = cute::cast_smem_ptr_to_uint(barrier); + asm volatile( + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes " + "[%0], [%1], %2, [%3];\n" + : + : "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(barrier_addr) + : "memory"); +#else + (void)smem_ptr; + (void)gmem_ptr; + (void)bytes; + (void)barrier; +#endif +} + CUTE_DEVICE void cp_async_commit_group() { asm volatile("cp.async.commit_group;\n" ::); } @@ -671,8 +694,8 @@ __global__ void qwen35_layout_scalar_kda_decode_long_kernel( } -template -__global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( +template +__global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( const scalar_t* __restrict__ mixed_qkv_conv, const scalar_t* __restrict__ a, const scalar_t* __restrict__ b, @@ -688,29 +711,34 @@ __global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( constexpr int kLocalKDim = Shape::kLocalKDim; constexpr int kLocalMixedQKVDim = Shape::kLocalMixedQKVDim; constexpr int kWarpSize = 32; - constexpr int kThreads = 32; - constexpr int kTileV = 32; + constexpr int kThreads = kTileV; + constexpr int kWarps = kThreads / kWarpSize; + constexpr int kVTiles = kHeadDimV / kTileV; constexpr int kKPerThread = kHeadDimQK / kThreads; - constexpr int kVecFloats = 4; static_assert(kLocalQKHeads == Shape::kLocalQKHeads); static_assert(kHeadDimQK == 128); static_assert(kHeadDimV == 128); + static_assert(kTileV == 32 || kTileV == 64); + static_assert(kHeadDimV % kTileV == 0); static_assert(kHeadDimQK % kThreads == 0); __shared__ float q_smem[kHeadDimQK]; __shared__ float k_smem[kHeadDimQK]; __shared__ float state_smem[kHeadDimQK][kTileV]; __shared__ float norm_smem[3]; + __shared__ float warp_reduce_smem[2 * kWarps]; + __shared__ cutlass::arch::ClusterTransactionBarrier::ValueType state_barrier; const int hv_tile = static_cast(blockIdx.x); const int token_idx = static_cast(blockIdx.y); - const int hv = hv_tile >> 2; - const int v_tile = hv_tile & 3; + const int hv = hv_tile / kVTiles; + const int v_tile = hv_tile - hv * kVTiles; const int tid = static_cast(threadIdx.x); const int lane = tid & (kWarpSize - 1); + const int warp_id = tid / kWarpSize; const int mapped_h = hv / kRepeatFactor; - const int v_row = v_tile * kTileV + lane; + const int v_row = v_tile * kTileV + tid; auto qk_src_layout = make_layout( make_shape(token_count, Int{}, Int{}), @@ -769,8 +797,19 @@ __global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( q_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_norm_sq); k_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_norm_sq); if (lane == 0) { - norm_smem[0] = rsqrtf(q_norm_sq + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); - norm_smem[1] = rsqrtf(k_norm_sq + 1e-6f); + warp_reduce_smem[warp_id] = q_norm_sq; + warp_reduce_smem[kWarps + warp_id] = k_norm_sq; + } + __syncthreads(); + if (warp_id == 0) { + float q_block_sum = lane < kWarps ? warp_reduce_smem[lane] : 0.f; + float k_block_sum = lane < kWarps ? warp_reduce_smem[kWarps + lane] : 0.f; + q_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_block_sum); + k_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_block_sum); + if (lane == 0) { + norm_smem[0] = rsqrtf(q_block_sum + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); + norm_smem[1] = rsqrtf(k_block_sum + 1e-6f); + } } __syncthreads(); @@ -786,7 +825,15 @@ __global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( } qk_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_dot); if (lane == 0) { - norm_smem[2] = qk_dot; + warp_reduce_smem[warp_id] = qk_dot; + } + __syncthreads(); + if (warp_id == 0) { + float qk_block_sum = lane < kWarps ? warp_reduce_smem[lane] : 0.f; + qk_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_block_sum); + if (lane == 0) { + norm_smem[2] = qk_block_sum; + } } __syncthreads(); @@ -797,8 +844,26 @@ __global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( const float decay = expf(g); const float beta = 1.f / (1.f + expf(-b_val)); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (tid == 0) { + cutlass::arch::ClusterTransactionBarrier::init(&state_barrier, 1); + cutlass::arch::ClusterTransactionBarrier::arrive_and_expect_tx( + &state_barrier, kHeadDimQK * kTileV * sizeof(float)); + } + __syncthreads(); +#pragma unroll 1 + for (int k_idx = tid; k_idx < kHeadDimQK; k_idx += kThreads) { + cp_async_bulk_shared_global( + &state_smem[k_idx][0], + &state_vk(v_tile * kTileV, k_idx), + kTileV * sizeof(float), + &state_barrier); + } + cutlass::arch::ClusterTransactionBarrier::wait(&state_barrier, 0); + __syncthreads(); +#else #pragma unroll 1 - for (int elem = tid * kVecFloats; elem < kHeadDimQK * kTileV; elem += kThreads * kVecFloats) { + for (int elem = tid * 4; elem < kHeadDimQK * kTileV; elem += kThreads * 4) { const int k_idx = elem / kTileV; const int v_base = elem - k_idx * kTileV; cp_async_ca_shared_global<16>( @@ -808,6 +873,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( cp_async_commit_group(); cp_async_wait_all(); __syncthreads(); +#endif float proj_acc0 = 0.f; float proj_acc1 = 0.f; @@ -827,14 +893,14 @@ __global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( float out_acc7 = 0.f; #pragma unroll 1 for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 8) { - const float state0 = state_smem[k_idx + 0][lane]; - const float state1 = state_smem[k_idx + 1][lane]; - const float state2 = state_smem[k_idx + 2][lane]; - const float state3 = state_smem[k_idx + 3][lane]; - const float state4 = state_smem[k_idx + 4][lane]; - const float state5 = state_smem[k_idx + 5][lane]; - const float state6 = state_smem[k_idx + 6][lane]; - const float state7 = state_smem[k_idx + 7][lane]; + const float state0 = state_smem[k_idx + 0][tid]; + const float state1 = state_smem[k_idx + 1][tid]; + const float state2 = state_smem[k_idx + 2][tid]; + const float state3 = state_smem[k_idx + 3][tid]; + const float state4 = state_smem[k_idx + 4][tid]; + const float state5 = state_smem[k_idx + 5][tid]; + const float state6 = state_smem[k_idx + 6][tid]; + const float state7 = state_smem[k_idx + 7][tid]; const float k0 = k_smem[k_idx + 0]; const float k1 = k_smem[k_idx + 1]; const float k2 = k_smem[k_idx + 2]; @@ -881,14 +947,14 @@ __global__ void qwen35_layout_scalar_kda_decode_long_v32_kernel( #pragma unroll 1 for (int k_idx = 0; k_idx < kHeadDimQK; k_idx += 8) { - const float state0 = state_smem[k_idx + 0][lane]; - const float state1 = state_smem[k_idx + 1][lane]; - const float state2 = state_smem[k_idx + 2][lane]; - const float state3 = state_smem[k_idx + 3][lane]; - const float state4 = state_smem[k_idx + 4][lane]; - const float state5 = state_smem[k_idx + 5][lane]; - const float state6 = state_smem[k_idx + 6][lane]; - const float state7 = state_smem[k_idx + 7][lane]; + const float state0 = state_smem[k_idx + 0][tid]; + const float state1 = state_smem[k_idx + 1][tid]; + const float state2 = state_smem[k_idx + 2][tid]; + const float state3 = state_smem[k_idx + 3][tid]; + const float state4 = state_smem[k_idx + 4][tid]; + const float state5 = state_smem[k_idx + 5][tid]; + const float state6 = state_smem[k_idx + 6][tid]; + const float state7 = state_smem[k_idx + 7][tid]; const float state_new0 = decay * state0 + v_new_row * k_smem[k_idx + 0]; const float state_new1 = decay * state1 + v_new_row * k_smem[k_idx + 1]; const float state_new2 = decay * state2 + v_new_row * k_smem[k_idx + 2]; @@ -924,9 +990,11 @@ void launch_qwen35_layout_scalar_kda_decode_long_kernel( constexpr int kWarpTileV = 32; (void)kWarpTileV; if (token_count == 64 || token_count == 128) { - dim3 grid(kLocalVHeads * 4, token_count, 1); - dim3 block(32, 1, 1); - qwen35_layout_scalar_kda_decode_long_v32_kernel<<>>( + constexpr int kLongTileV = 64; + dim3 grid(kLocalVHeads * (kHeadDimV / kLongTileV), token_count, 1); + dim3 block(kLongTileV, 1, 1); + qwen35_layout_scalar_kda_decode_long_vtile_kernel + <<>>( mixed_qkv_conv, a, b, diff --git a/tests/test_qwen35_decode.py b/tests/test_qwen35_decode.py index 274ff85d..7f206185 100644 --- a/tests/test_qwen35_decode.py +++ b/tests/test_qwen35_decode.py @@ -526,6 +526,44 @@ def test_qwen35_fused_layout_kda_cudac_matches_reference_unfused_and_triton(toke assert torch.allclose(state_triton, state_fused, atol=3e-5, rtol=3e-5) +@pytest.mark.skipif( + not _has_qwen35_fused_layout_kda_cudac(), + reason="Qwen3.5 fused layout+KDA CUDA backend is not available", +) +@pytest.mark.parametrize("tokens", [64, 128]) +def test_qwen35_fused_layout_kda_cudac_long_matches_reference(tokens: int): + mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = make_inputs( + tokens=tokens, + pool_size=tokens, + device=torch.device("cuda"), + ) + conv_out, _ = qwen35_conv1d_decode_update( + mixed_qkv, + conv_state, + conv_weight, + activation="silu", + backend="cudac", + ) + out_ref, state_ref = manual_qwen35_layout_scalar_kda_reference( + conv_out, a, b, A_log, dt_bias, recurrent_state, state_indices + ) + out_fused, state_fused = qwen35_layout_scalar_kda_decode( + mixed_qkv_conv=conv_out, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + state_indices=state_indices, + config=DEFAULT_QWEN35_LINEAR_ATTN_CONFIG, + backend="cudac", + ) + + torch.cuda.synchronize() + torch.testing.assert_close(out_fused.float(), out_ref.float(), atol=3e-2, rtol=3e-2) + torch.testing.assert_close(state_fused, state_ref, atol=3e-5, rtol=3e-5) + + @pytest.mark.skipif(not _has_qwen35_fused_layout_kda_cudac(), reason="Qwen3.5 fused layout+KDA CUDA backend is not available") @pytest.mark.parametrize("local_v_heads", [48, 24, 12, 6]) def test_qwen35_layout_scalar_kda_cudac_supports_local_tp_shards(local_v_heads: int): From 0c2450c63d243d717c858cbb0253bcdb6cc5d27a Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sun, 12 Jul 2026 09:19:00 +0000 Subject: [PATCH 19/35] feat(kda): support native GVA prefill benchmarks --- benchmarks/bench_kda_fused_fwd.py | 34 +- benchmarks/bench_qwen35_decode.py | 739 +++++--------------------- benchmarks/bench_qwen35_prefill.py | 449 ---------------- cula/kda/blackwell_fused_fwd.py | 10 +- cula/ops/kda_fully_fused_sm100_wip.py | 57 +- 5 files changed, 209 insertions(+), 1080 deletions(-) delete mode 100644 benchmarks/bench_qwen35_prefill.py diff --git a/benchmarks/bench_kda_fused_fwd.py b/benchmarks/bench_kda_fused_fwd.py index b42a0f7e..15296617 100644 --- a/benchmarks/bench_kda_fused_fwd.py +++ b/benchmarks/bench_kda_fused_fwd.py @@ -176,13 +176,22 @@ def bench_fixed(configs): cu_seqlens=cu_seqlens, lower_bound=lower_bound, ) + common_cula = dict(common) + if init_state is not None: + common_cula["init_state"] = init_state.transpose(-1, -2).contiguous() # Accuracy - o_fla, _ = run_fla(**common) - o_cula, _ = run_cula(**common) + o_fla, ht_fla = run_fla(**common) + o_cula, ht_cula_vk = run_cula(**common_cula) torch.cuda.synchronize() relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) + if ht_fla is not None and ht_cula_vk is not None: + ht_cula = ht_cula_vk.transpose(-1, -2) + state_rms, state_max, state_mean = relative_rms_error_rel_max_mean_abs(ht_fla, ht_cula) + relative_rms_error = max(relative_rms_error, state_rms) + rel_max = max(rel_max, state_max) + mean_diff = max(mean_diff, state_mean) # Performance ms_fla = benchmark_cuda_mode_fn( @@ -193,7 +202,7 @@ def bench_fixed(configs): sanitizer_mode=SANITIZER_MODE, ) ms_cula = benchmark_cuda_mode_fn( - lambda: run_cula(**common), + lambda: run_cula(**common_cula), default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE, @@ -216,7 +225,7 @@ def bench_fixed(configs): } ) - del o_fla, o_cula, q, k, v, g, beta, A_log, dt_bias, inputs + del o_fla, o_cula, ht_fla, ht_cula_vk, q, k, v, g, beta, A_log, dt_bias, inputs torch.cuda.empty_cache() return results @@ -267,13 +276,22 @@ def bench_varlen(configs): cu_seqlens=cu_seqlens, lower_bound=lower_bound, ) + common_cula = dict(common) + if init_state is not None: + common_cula["init_state"] = init_state.transpose(-1, -2).contiguous() # Accuracy - o_fla, _ = run_fla(**common) - o_cula, _ = run_cula(**common) + o_fla, ht_fla = run_fla(**common) + o_cula, ht_cula_vk = run_cula(**common_cula) torch.cuda.synchronize() relative_rms_error, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) + if ht_fla is not None and ht_cula_vk is not None: + ht_cula = ht_cula_vk.transpose(-1, -2) + state_rms, state_max, state_mean = relative_rms_error_rel_max_mean_abs(ht_fla, ht_cula) + relative_rms_error = max(relative_rms_error, state_rms) + rel_max = max(rel_max, state_max) + mean_diff = max(mean_diff, state_mean) # Performance ms_fla = benchmark_cuda_mode_fn( @@ -284,7 +302,7 @@ def bench_varlen(configs): sanitizer_mode=SANITIZER_MODE, ) ms_cula = benchmark_cuda_mode_fn( - lambda: run_cula(**common), + lambda: run_cula(**common_cula), default_warmup=WARMUP, default_rep=N_ITERS, ncu_mode=NCU_MODE, @@ -314,7 +332,7 @@ def bench_varlen(configs): } ) - del o_fla, o_cula, q, k, v, g, beta, A_log, dt_bias, inputs + del o_fla, o_cula, ht_fla, ht_cula_vk, q, k, v, g, beta, A_log, dt_bias, inputs torch.cuda.empty_cache() return results diff --git a/benchmarks/bench_qwen35_decode.py b/benchmarks/bench_qwen35_decode.py index 7cd9f858..dbf4cbc2 100755 --- a/benchmarks/bench_qwen35_decode.py +++ b/benchmarks/bench_qwen35_decode.py @@ -1,667 +1,208 @@ #!/usr/bin/env python3 -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Benchmark Qwen3.5 decode on the active CUDA device. - -Two timing scopes are reported: - - native_core: direct native scalar GDN decode op. - - triton_core: FLA/SGLang-style fused_sigmoid_gating_delta_rule_update - Triton decode op vendored in cuLA. - - sglang_core: fused_sigmoid_gating_delta_rule_update from SGLang, when - available from the installed package or --sglang-path. - - fused_layout_kda: direct cuLA fused Qwen3.5 layout + scalar KDA decode op. - - sglang_packed: SGLang packed Qwen3.5 layout + recurrent update op, when - available from the installed package or --sglang-path. - - full: cuLA Python Qwen3.5 decode chain, including conv + layout + core. - -State buffers are reset before each timed iteration and the reset copy is not -included in the event timing window. +"""Benchmark the fused cuLA Qwen3.5 decode kernel against upstream FLA. + +The upstream FLA recurrent operator receives pre-laid-out Q/K/V tensors. cuLA +receives the packed Qwen3.5 ``mixed_qkv_conv`` tensor and performs layout plus +the recurrent update in one kernel. State reset is outside both timing windows. """ from __future__ import annotations import argparse import csv -import importlib -import importlib.util -import inspect import pathlib import statistics import sys -import time -from collections.abc import Callable -import torch +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -ROOT = pathlib.Path(__file__).resolve().parent.parent -sys.path.insert(0, str(ROOT)) +import torch +from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule import cula.cudac as cula_cuda -from cula.ops.kda_decode_fla import fused_sigmoid_gating_delta_rule_update as triton_fused_sigmoid_update -from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as CONFIG +from cula.ops.qwen35_layout_decode import qwen35_layout_decode_reference +from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as GLOBAL_CONFIG from cula.qwen35.common import Qwen35LinearAttentionConfig -from cula.qwen35.runtime import qwen35_linear_attention_decode - -SGLANG_CORE_MODULES = [ - "sglang.srt.layers.attention.linear.kernels.gdn_triton", - "sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent", -] -SGLANG_CORE_FILES = [ - pathlib.Path("sglang/srt/layers/attention/linear/kernels/gdn_triton.py"), - pathlib.Path("sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py"), -] -SGLANG_PACKED_MODULES = [ - "sglang.srt.layers.attention.fla.fused_recurrent", - "sglang.srt.layers.attention.linear.kernels.gdn_triton", -] -SGLANG_PACKED_FILES = [ - pathlib.Path("sglang/srt/layers/attention/fla/fused_recurrent.py"), - pathlib.Path("sglang/srt/layers/attention/linear/kernels/gdn_triton.py"), -] - - -def accelerator_device() -> torch.device: - if torch.cuda.is_available(): - return torch.device("cuda") - raise RuntimeError("No CUDA accelerator is available.") - - -def accelerator_name(device: torch.device) -> str: - if device.type != "cuda": - raise ValueError(f"Unsupported device={device}") - return torch.cuda.get_device_name(device.index or 0) - - -def synchronize(device: torch.device) -> None: - if device.type != "cuda": - raise ValueError(f"Unsupported device={device}") - torch.cuda.synchronize() -def benchmark_accel_fn( - fn: Callable[[], object], - *, - device: torch.device, - setup_fn: Callable[[], None] | None, - warmup: int, - rep: int, -) -> float: +def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: + return Qwen35LinearAttentionConfig( + hidden_size=GLOBAL_CONFIG.hidden_size // tp_size, + conv_kernel_size=GLOBAL_CONFIG.conv_kernel_size, + num_k_heads=GLOBAL_CONFIG.num_k_heads // tp_size, + num_v_heads=GLOBAL_CONFIG.num_v_heads // tp_size, + head_k_dim=GLOBAL_CONFIG.head_k_dim, + head_v_dim=GLOBAL_CONFIG.head_v_dim, + qkv_dtype=GLOBAL_CONFIG.qkv_dtype, + state_dtype=GLOBAL_CONFIG.state_dtype, + ) + + +def benchmark_cuda(fn, *, setup=None, warmup: int, rep: int) -> float: for _ in range(warmup): - if setup_fn is not None: - setup_fn() + if setup is not None: + setup() fn() - synchronize(device) - - times: list[float] = [] - try: - starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - for i in range(rep): - if setup_fn is not None: - setup_fn() - starts[i].record() - fn() - ends[i].record() - synchronize(device) - times = [s.elapsed_time(e) for s, e in zip(starts, ends)] - except Exception: - for _ in range(rep): - if setup_fn is not None: - setup_fn() - synchronize(device) - t0 = time.perf_counter() - fn() - synchronize(device) - times.append((time.perf_counter() - t0) * 1000.0) - - if not times: - return 0.0 - if len(times) < 4: - return statistics.mean(times) - times = sorted(times) - iqr = times[len(times) // 4 : 3 * len(times) // 4] - return statistics.mean(iqr) + torch.cuda.synchronize() + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + for i in range(rep): + if setup is not None: + setup() + starts[i].record() + fn() + ends[i].record() + torch.cuda.synchronize() + + times = sorted(start.elapsed_time(end) for start, end in zip(starts, ends)) + lo, hi = len(times) // 4, 3 * len(times) // 4 + return statistics.mean(times[lo:hi] or times) -def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: - if tp_size not in (1, 2, 4, 8): - raise ValueError(f"tp_size must be one of 1, 2, 4, 8, got {tp_size}") - return Qwen35LinearAttentionConfig( - hidden_size=CONFIG.hidden_size // tp_size, - conv_kernel_size=CONFIG.conv_kernel_size, - num_k_heads=CONFIG.num_k_heads // tp_size, - num_v_heads=CONFIG.num_v_heads // tp_size, - head_k_dim=CONFIG.head_k_dim, - head_v_dim=CONFIG.head_v_dim, - qkv_dtype=CONFIG.qkv_dtype, - state_dtype=CONFIG.state_dtype, - ) +def error_stats(reference: torch.Tensor, actual: torch.Tensor) -> tuple[float, float]: + reference = reference.float() + actual = actual.float() + diff = (reference - actual).abs() + rel_rms = diff.square().mean().sqrt() / reference.square().mean().sqrt().clamp_min(1e-8) + return rel_rms.item(), diff.max().item() -def make_full_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): - torch.manual_seed(seed) - pool_size = max(tokens, 1) - mixed_qkv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) - a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) - b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) - conv_weight = torch.randn(config.conv_dim, config.conv_kernel_size, device=device, dtype=config.qkv_dtype) - conv_state = torch.randn( + +def make_inputs(tokens: int, *, tp_size: int, seed: int, device: torch.device): + config = local_config_from_tp_size(tp_size) + generator = torch.Generator(device=device).manual_seed(seed) + hv, k_dim, v_dim = config.num_v_heads, config.head_k_dim, config.head_v_dim + + mixed_qkv = torch.randn( tokens, config.conv_dim, - config.conv_kernel_size, + generator=generator, device=device, dtype=config.qkv_dtype, ) - recurrent_state = torch.randn( - pool_size, - config.num_v_heads, - config.head_k_dim, - config.head_v_dim, - device=device, - dtype=config.state_dtype, - ) * 0.01 - A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 - state_indices = torch.arange(tokens, device=device, dtype=torch.int32) - return mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices - - -def make_fused_layout_kda_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): - torch.manual_seed(seed) - mixed_qkv_conv = torch.randn(tokens, config.conv_dim, device=device, dtype=config.qkv_dtype) - a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) - b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) - A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 + a = torch.randn(tokens, hv, generator=generator, device=device, dtype=config.qkv_dtype) + b = torch.randn(tokens, hv, generator=generator, device=device, dtype=config.qkv_dtype) + A_log = -torch.rand(hv, generator=generator, device=device, dtype=torch.float32) + dt_bias = torch.randn(hv, generator=generator, device=device, dtype=torch.float32) * 0.1 state = torch.randn( tokens, - config.num_v_heads, - config.head_k_dim, - config.head_v_dim, + hv, + k_dim, + v_dim, + generator=generator, device=device, - dtype=config.state_dtype, + dtype=torch.float32, ) * 0.01 - state_work = torch.empty_like(state) - state_indices = torch.arange(tokens, device=device, dtype=torch.int32) - out = torch.empty(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) - return mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out - - -def make_core_inputs(tokens: int, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): - torch.manual_seed(seed) - q = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) - k = torch.randn(tokens, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) - v = torch.randn(tokens, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) - a = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) - b = torch.randn(tokens, config.num_v_heads, device=device, dtype=config.qkv_dtype) - A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 - state = torch.randn( - tokens, - config.num_v_heads, - config.head_k_dim, - config.head_v_dim, - device=device, - dtype=config.state_dtype, - ) * 0.01 - state_indices = torch.arange(tokens, device=device, dtype=torch.int32) - out = torch.empty_like(v) - state_work = torch.empty_like(state) - return q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out - - -def _add_sglang_import_roots(sglang_path: pathlib.Path | None) -> None: - if sglang_path is not None: - for import_root in (sglang_path, sglang_path / "python"): - if import_root.exists(): - sys.path.insert(0, str(import_root)) - - -def _resolve_sglang_symbol( - *, - sglang_path: pathlib.Path | None, - module_names: list[str], - file_paths: list[pathlib.Path], - symbol_names: list[str], -): - _add_sglang_import_roots(sglang_path) - - import_errors = [] - for module_name in module_names: - try: - module = importlib.import_module(module_name) - except (ImportError, PermissionError, ModuleNotFoundError) as exc: - import_errors.append(f"{module_name}: {type(exc).__name__}: {exc}") - continue - for symbol_name in symbol_names: - if hasattr(module, symbol_name): - return getattr(module, symbol_name), f"{module_name}.{symbol_name}" - - if sglang_path is not None: - candidates: list[pathlib.Path] = [] - for rel_path in file_paths: - candidates.extend([sglang_path / rel_path, sglang_path / "python" / rel_path]) - for idx, path in enumerate(candidates): - if path.exists(): - spec = importlib.util.spec_from_file_location(f"_sglang_qwen35_decode_provider_{idx}", path) - if spec is None or spec.loader is None: - continue - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - for symbol_name in symbol_names: - if hasattr(module, symbol_name): - return getattr(module, symbol_name), f"{path}:{symbol_name}" - raise RuntimeError( - f"Could not find any of {symbol_names} under --sglang-path={sglang_path}. " - "Pass the SGLang repo root or its python/ directory. " - f"Import errors: {'; '.join(import_errors) or 'none'}" - ) - - return None, None - - -def resolve_sglang_core_update(sglang_path: pathlib.Path | None): - """Return SGLang's scalar-gated recurrent update function if available.""" - return _resolve_sglang_symbol( - sglang_path=sglang_path, - module_names=SGLANG_CORE_MODULES, - file_paths=SGLANG_CORE_FILES, - symbol_names=["fused_sigmoid_gating_delta_rule_update"], - ) - - -def resolve_sglang_packed_decode(sglang_path: pathlib.Path | None): - """Return SGLang's packed layout + recurrent decode function if available.""" - return _resolve_sglang_symbol( - sglang_path=sglang_path, - module_names=SGLANG_PACKED_MODULES, - file_paths=SGLANG_PACKED_FILES, - symbol_names=[ - "fused_recurrent_gated_delta_rule_packed_decode", - "fused_recurrent_gated_delta_rule_packed_decode_cpu", - ], - ) - - -def call_with_supported_kwargs(fn: Callable, **kwargs): - """Call a provider while tolerating minor SGLang signature drift.""" - try: - signature = inspect.signature(fn) - except (TypeError, ValueError): - return fn(**kwargs) - if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): - return fn(**kwargs) - filtered = {name: value for name, value in kwargs.items() if name in signature.parameters} - return fn(**filtered) - - -def bench_native_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: - q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_core_inputs(tokens, device, seed, config) + indices = torch.arange(tokens, device=device, dtype=torch.int32) + return config, mixed_qkv, a, b, A_log, dt_bias, state, indices - def setup() -> None: - state_work.copy_(state) - def run() -> None: - cula_cuda.qwen35_scalar_kda_decode( - q, - k, - v, - a, - b, - A_log, - dt_bias, - state_work, - state_indices, - out, - ) - - return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) - - -def bench_triton_core(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: - q, k, v, a, b, A_log, dt_bias, state, state_work, state_indices, _ = make_core_inputs(tokens, device, seed, config) - q_4d = q.unsqueeze(1).contiguous() - k_4d = k.unsqueeze(1).contiguous() - v_4d = v.unsqueeze(1).contiguous() - a_3d = a.unsqueeze(1).contiguous() - b_3d = b.unsqueeze(1).contiguous() - - def setup() -> None: - state_work.copy_(state) - - def run() -> None: - triton_fused_sigmoid_update( - A_log=A_log, - a=a_3d, - dt_bias=dt_bias, - softplus_beta=1.0, - softplus_threshold=20.0, - q=q_4d, - k=k_4d, - v=v_4d, - b=b_3d, - initial_state_source=state_work, - initial_state_indices=state_indices, - scale=config.head_k_dim**-0.5, - use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - is_kda=False, - ) - - return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) - - -def bench_fused_layout_kda(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: - mixed_qkv_conv, a, b, A_log, dt_bias, state, state_work, state_indices, out = make_fused_layout_kda_inputs( - tokens, device, seed, config +def run_case(tokens: int, args, device: torch.device) -> dict[str, float | int]: + config, mixed, a, b, A_log, dt_bias, state, indices = make_inputs( + tokens, + tp_size=args.tp_size, + seed=args.seed, + device=device, ) - - def setup() -> None: - state_work.copy_(state) - - def run() -> None: - cula_cuda.qwen35_layout_scalar_kda_decode( - mixed_qkv_conv, - a, - b, - A_log, - dt_bias, - state_work, - state_indices, - out, - ) - - return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) - - -def bench_sglang_core( - tokens: int, - device: torch.device, - warmup: int, - rep: int, - seed: int, - sglang_fused_update: Callable, - config: Qwen35LinearAttentionConfig, -) -> float: - q, k, v, a, b, A_log, dt_bias, state, _, state_indices, _ = make_core_inputs(tokens, device, seed, config) - q_4d = q.unsqueeze(1).contiguous() - k_4d = k.unsqueeze(1).contiguous() - v_4d = v.unsqueeze(1).contiguous() - a_3d = a.unsqueeze(1).contiguous() - b_3d = b.unsqueeze(1).contiguous() - state_vk = state.transpose(-1, -2).contiguous() - state_vk_work = torch.empty_like(state_vk) - - def setup() -> None: - state_vk_work.copy_(state_vk) - - def run() -> None: - call_with_supported_kwargs( - sglang_fused_update, - A_log=A_log, - a=a_3d, - dt_bias=dt_bias, - softplus_beta=1.0, - softplus_threshold=20.0, - q=q_4d, - k=k_4d, - v=v_4d, - b=b_3d, - initial_state_source=state_vk_work, - initial_state_indices=state_indices, + q, k, v, a_fla, b_fla = qwen35_layout_decode_reference(mixed, a, b, config=config) + q, k, v = q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1) + gate = a_fla.unsqueeze(1) + beta = torch.sigmoid(b_fla.float()).unsqueeze(1) + + state_cula = torch.empty_like(state) + out_cula = torch.empty_like(v.squeeze(1)) + + def run_fla(): + return fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=gate, + beta=beta, scale=config.head_k_dim**-0.5, + initial_state=state, + output_final_state=True, use_qk_l2norm_in_kernel=True, - cu_seqlens=None, - is_kda=False, - ) - - return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) - - -def bench_sglang_packed_layout_kda( - tokens: int, - device: torch.device, - warmup: int, - rep: int, - seed: int, - sglang_packed_decode: Callable, - config: Qwen35LinearAttentionConfig, -) -> float: - mixed_qkv_conv, a, b, A_log, dt_bias, state, _, state_indices, _ = make_fused_layout_kda_inputs(tokens, device, seed, config) - state_vk = state.transpose(-1, -2).contiguous() - state_vk_work = torch.empty_like(state_vk) - out = torch.empty(tokens, 1, config.num_v_heads, config.head_v_dim, device=device, dtype=config.qkv_dtype) - - def setup() -> None: - state_vk_work.copy_(state_vk) - - def run() -> None: - call_with_supported_kwargs( - sglang_packed_decode, - mixed_qkv=mixed_qkv_conv, - a=a, - b=b, + use_gate_in_kernel=True, A_log=A_log, dt_bias=dt_bias, - scale=config.head_k_dim**-0.5, - initial_state=state_vk_work, - out=out, - ssm_state_indices=state_indices, - use_qk_l2norm_in_kernel=True, + transpose_state_layout=False, ) - return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) + def setup_cula() -> None: + state_cula.copy_(state) - -def bench_full(tokens: int, device: torch.device, warmup: int, rep: int, seed: int, config: Qwen35LinearAttentionConfig) -> float: - inputs = make_full_inputs(tokens, device, seed, config) - mixed_qkv, a, b, conv_weight, conv_state, recurrent_state, A_log, dt_bias, state_indices = inputs - conv_state_work = torch.empty_like(conv_state) - recurrent_state_work = torch.empty_like(recurrent_state) - - def setup() -> None: - conv_state_work.copy_(conv_state) - recurrent_state_work.copy_(recurrent_state) - - def run() -> None: - qwen35_linear_attention_decode( - mixed_qkv, + def run_cula() -> None: + cula_cuda.qwen35_layout_scalar_kda_decode( + mixed, a, b, - conv_weight, A_log, dt_bias, - config=config, - conv_state=conv_state_work, - recurrent_state=recurrent_state_work, - state_indices=state_indices, - backend="cudac", + state_cula, + indices, + out_cula, ) - return benchmark_accel_fn(run, device=device, setup_fn=setup, warmup=warmup, rep=rep) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Benchmark cuLA Qwen3.5 decode.") + out_fla, state_fla = run_fla() + setup_cula() + run_cula() + torch.cuda.synchronize() + out_rel_rms, out_max_abs = error_stats(out_fla.squeeze(1), out_cula) + state_rel_rms, state_max_abs = error_stats(state_fla, state_cula) + + fla_ms = benchmark_cuda(run_fla, warmup=args.warmup, rep=args.rep) + cula_ms = benchmark_cuda(run_cula, setup=setup_cula, warmup=args.warmup, rep=args.rep) + return { + "tokens": tokens, + "upstream_fla_ms": fla_ms, + "cula_fused_ms": cula_ms, + "speedup": fla_ms / cula_ms, + "out_rel_rms": out_rel_rms, + "out_max_abs": out_max_abs, + "state_rel_rms": state_rel_rms, + "state_max_abs": state_max_abs, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Upstream FLA vs fused cuLA Qwen3.5 decode") parser.add_argument("--tokens", nargs="+", type=int, default=[1, 2, 4, 8, 16, 32, 64, 128]) parser.add_argument("--warmup", type=int, default=20) parser.add_argument("--rep", type=int, default=100) parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--scope", choices=["core", "fused", "full", "both"], default="both") parser.add_argument("--tp-size", type=int, choices=[1, 2, 4, 8], default=1) - parser.add_argument("--skip-triton", action="store_true", help="Skip the vendored Triton core timing.") - parser.add_argument("--skip-sglang", action="store_true", help="Do not try the SGLang kernel provider.") - parser.add_argument("--require-sglang", action="store_true", help="Fail if the SGLang kernel provider is unavailable.") - parser.add_argument("--sglang-path", type=pathlib.Path, default=None, help="SGLang repo root or python/ directory.") - parser.add_argument("--csv", type=pathlib.Path, default=None) - return parser.parse_args() - - -def main() -> int: - args = parse_args() - config = local_config_from_tp_size(args.tp_size) - device = accelerator_device() - rows: list[dict[str, object]] = [] - sglang_fused_update = None - sglang_core_source = None - sglang_packed_decode = None - sglang_packed_source = None - if not args.skip_sglang: - sglang_fused_update, sglang_core_source = resolve_sglang_core_update(args.sglang_path) - sglang_packed_decode, sglang_packed_source = resolve_sglang_packed_decode(args.sglang_path) - if args.require_sglang and (sglang_fused_update is None or sglang_packed_decode is None): - raise RuntimeError("SGLang core and packed decode providers must both be available.") - - print(f"device={device} name={accelerator_name(device)} torch={torch.__version__}") + parser.add_argument("--csv", type=pathlib.Path) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + device = torch.device("cuda") print( - f"qwen35: tp={args.tp_size} local_HK={config.num_k_heads} local_HV={config.num_v_heads} " - f"K={config.head_k_dim} V={config.head_v_dim} conv_dim={config.conv_dim}" + f"device={torch.cuda.get_device_name(device)} torch={torch.__version__} " + f"cuda={torch.version.cuda} tp={args.tp_size} warmup/rep={args.warmup}/{args.rep}" ) - print(f"sglang_core_provider={sglang_core_source or 'unavailable'}") - print(f"sglang_packed_provider={sglang_packed_source or 'unavailable'}") - print("| tokens | native_core_ms | triton_core_ms | sglang_core_ms | fused_layout_kda_ms | sglang_packed_ms | full_ms | triton/native | sglang/native | packed/fused | native_us_per_token | triton_us_per_token | sglang_us_per_token | fused_us_per_token | packed_us_per_token | full_us_per_token |") - print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|") + print("scope: upstream FLA recurrent operator vs cuLA fused Qwen layout + recurrent kernel") + print("| tokens | upstream_fla_ms | cula_fused_ms | speedup | out_rel_rms | state_rel_rms |") + print("|---:|---:|---:|---:|---:|---:|") + rows = [] for tokens in args.tokens: - native_core_ms = None - triton_core_ms = None - sglang_core_ms = None - fused_layout_kda_ms = None - sglang_packed_ms = None - full_ms = None - if args.scope in ("core", "both"): - native_core_ms = bench_native_core(tokens, device, args.warmup, args.rep, args.seed, config) - if not args.skip_triton: - triton_core_ms = bench_triton_core(tokens, device, args.warmup, args.rep, args.seed, config) - if sglang_fused_update is not None: - sglang_core_ms = bench_sglang_core( - tokens, - device, - args.warmup, - args.rep, - args.seed, - sglang_fused_update, - config, - ) - if args.scope in ("fused", "both"): - fused_layout_kda_ms = bench_fused_layout_kda(tokens, device, args.warmup, args.rep, args.seed, config) - if sglang_packed_decode is not None: - sglang_packed_ms = bench_sglang_packed_layout_kda( - tokens, - device, - args.warmup, - args.rep, - args.seed, - sglang_packed_decode, - config, - ) - if args.scope in ("full", "both"): - full_ms = bench_full(tokens, device, args.warmup, args.rep, args.seed, config) - - native_core_us = None if native_core_ms is None else native_core_ms * 1000.0 / tokens - triton_core_us = None if triton_core_ms is None else triton_core_ms * 1000.0 / tokens - sglang_core_us = None if sglang_core_ms is None else sglang_core_ms * 1000.0 / tokens - fused_layout_kda_us = None if fused_layout_kda_ms is None else fused_layout_kda_ms * 1000.0 / tokens - sglang_packed_us = None if sglang_packed_ms is None else sglang_packed_ms * 1000.0 / tokens - full_us = None if full_ms is None else full_ms * 1000.0 / tokens - triton_ratio = None - if native_core_ms is not None and triton_core_ms is not None and native_core_ms > 0: - triton_ratio = triton_core_ms / native_core_ms - sglang_ratio = None - if native_core_ms is not None and sglang_core_ms is not None and native_core_ms > 0: - sglang_ratio = sglang_core_ms / native_core_ms - packed_ratio = None - if fused_layout_kda_ms is not None and sglang_packed_ms is not None and fused_layout_kda_ms > 0: - packed_ratio = sglang_packed_ms / fused_layout_kda_ms + row = run_case(tokens, args, device) + rows.append(row) print( - f"| {tokens} | " - f"{'n/a' if native_core_ms is None else f'{native_core_ms:.4f}'} | " - f"{'n/a' if triton_core_ms is None else f'{triton_core_ms:.4f}'} | " - f"{'n/a' if sglang_core_ms is None else f'{sglang_core_ms:.4f}'} | " - f"{'n/a' if fused_layout_kda_ms is None else f'{fused_layout_kda_ms:.4f}'} | " - f"{'n/a' if sglang_packed_ms is None else f'{sglang_packed_ms:.4f}'} | " - f"{'n/a' if full_ms is None else f'{full_ms:.4f}'} | " - f"{'n/a' if triton_ratio is None else f'{triton_ratio:.2f}x'} | " - f"{'n/a' if sglang_ratio is None else f'{sglang_ratio:.2f}x'} | " - f"{'n/a' if packed_ratio is None else f'{packed_ratio:.2f}x'} | " - f"{'n/a' if native_core_us is None else f'{native_core_us:.2f}'} | " - f"{'n/a' if triton_core_us is None else f'{triton_core_us:.2f}'} | " - f"{'n/a' if sglang_core_us is None else f'{sglang_core_us:.2f}'} | " - f"{'n/a' if fused_layout_kda_us is None else f'{fused_layout_kda_us:.2f}'} | " - f"{'n/a' if sglang_packed_us is None else f'{sglang_packed_us:.2f}'} | " - f"{'n/a' if full_us is None else f'{full_us:.2f}'} |" - ) - rows.append( - { - "tokens": tokens, - "tp_size": args.tp_size, - "local_k_heads": config.num_k_heads, - "local_v_heads": config.num_v_heads, - "conv_dim": config.conv_dim, - "native_core_ms": native_core_ms, - "triton_core_ms": triton_core_ms, - "sglang_core_ms": sglang_core_ms, - "fused_layout_kda_ms": fused_layout_kda_ms, - "sglang_packed_ms": sglang_packed_ms, - "full_ms": full_ms, - "triton_over_native": triton_ratio, - "sglang_over_native": sglang_ratio, - "sglang_packed_over_fused": packed_ratio, - "native_core_us_per_token": native_core_us, - "triton_core_us_per_token": triton_core_us, - "sglang_core_us_per_token": sglang_core_us, - "fused_layout_kda_us_per_token": fused_layout_kda_us, - "sglang_packed_us_per_token": sglang_packed_us, - "full_us_per_token": full_us, - } + f"| {tokens} | {row['upstream_fla_ms']:.4f} | {row['cula_fused_ms']:.4f} | " + f"{row['speedup']:.2f}x | {row['out_rel_rms']:.3e} | {row['state_rel_rms']:.3e} |" ) if args.csv is not None: args.csv.parent.mkdir(parents=True, exist_ok=True) - with args.csv.open("w", newline="", encoding="utf-8") as f: - writer = csv.DictWriter( - f, - fieldnames=[ - "tokens", - "tp_size", - "local_k_heads", - "local_v_heads", - "conv_dim", - "native_core_ms", - "triton_core_ms", - "sglang_core_ms", - "fused_layout_kda_ms", - "sglang_packed_ms", - "full_ms", - "triton_over_native", - "sglang_over_native", - "sglang_packed_over_fused", - "native_core_us_per_token", - "triton_core_us_per_token", - "sglang_core_us_per_token", - "fused_layout_kda_us_per_token", - "sglang_packed_us_per_token", - "full_us_per_token", - ], - ) + with args.csv.open("w", newline="", encoding="utf-8") as file: + writer = csv.DictWriter(file, fieldnames=list(rows[0])) writer.writeheader() writer.writerows(rows) print(f"wrote {args.csv}") - return 0 - if __name__ == "__main__": - raise SystemExit(main()) + main() diff --git a/benchmarks/bench_qwen35_prefill.py b/benchmarks/bench_qwen35_prefill.py deleted file mode 100644 index bfcbe6d5..00000000 --- a/benchmarks/bench_qwen35_prefill.py +++ /dev/null @@ -1,449 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Benchmark Qwen3.5 prefill kernels. - -Reports: - - layout: cuLA Qwen3.5 prefill layout split/repeat kernel - - cula_qk: cuLA Qwen3.5 TMA/WGMMA-or-UMMA QK chunk debug kernel - - cula_fused: cuLA generic fused KDA core through a Qwen3.5 scalar-gate adapter - - fla_gdr: optional FLA chunk_gated_delta_rule baseline - - sgl_gdr: optional SGLang vendored Triton chunk_gated_delta_rule baseline - -Baselines are optional. SGLang Qwen3.5 prefill uses the same chunked gated -delta rule family in its Triton GDN kernel; decode uses a recurrent packed -kernel instead. - -Note: cula_qk currently benchmarks the TMA tensor-core Q @ K^T subpath only, -not the full gated-delta prefill recurrence. Its output is [B,local_HV,T,T], -so long sequence lengths have quadratic memory cost. -""" - -from __future__ import annotations - -import argparse -import importlib -import inspect -import pathlib -import statistics -import sys -from collections.abc import Callable - -import torch -import torch.nn.functional as F - -ROOT = pathlib.Path(__file__).resolve().parent.parent -sys.path.insert(0, str(ROOT)) - -from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill -from cula.ops.qwen35_fused_kda_prefill import qwen35_fused_kda_prefill -from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill -from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as CONFIG -from cula.qwen35.common import Qwen35LinearAttentionConfig -from cula.utils import get_kda_fused_fwd - -try: - import cula.cudac as cula_cuda -except ImportError: - cula_cuda = None - -RCP_LN2 = 1.4426950408889634 - - -def benchmark_cuda_fn(fn: Callable[[], object], *, warmup: int, rep: int) -> float: - for _ in range(warmup): - fn() - torch.cuda.synchronize() - - starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - for idx in range(rep): - starts[idx].record() - fn() - ends[idx].record() - torch.cuda.synchronize() - - times = [start.elapsed_time(end) for start, end in zip(starts, ends)] - if len(times) <= 2: - return statistics.mean(times) - times = sorted(times) - return statistics.mean(times[len(times) // 4 : 3 * len(times) // 4]) - - -def error_stats(ref: torch.Tensor, out: torch.Tensor) -> tuple[float, float, float]: - ref_f = ref.float() - out_f = out.float() - diff = (ref_f - out_f).abs() - rmse = diff.square().mean().sqrt().item() - ref_rms = ref_f.square().mean().sqrt().item() - rel_rms = rmse / (ref_rms + 1.0e-8) - rel_max = diff.max().item() / (ref_f.abs().max().item() + 1.0e-8) - mean_abs = diff.mean().item() - return rel_rms, rel_max, mean_abs - - -def resolve_fla_chunk_gdr(): - try: - module = importlib.import_module("fla.ops.gated_delta_rule") - except ImportError as exc: - return None, f"cannot import fla.ops.gated_delta_rule: {exc}" - if not hasattr(module, "chunk_gated_delta_rule"): - return None, "fla.ops.gated_delta_rule has no chunk_gated_delta_rule" - return module.chunk_gated_delta_rule, "fla.ops.gated_delta_rule.chunk_gated_delta_rule" - - -def resolve_sgl_chunk_gdr(sglang_path: pathlib.Path | None): - if sglang_path is not None: - for root in (sglang_path, sglang_path / "python"): - if root.exists(): - sys.path.insert(0, str(root)) - try: - module = importlib.import_module("sglang.srt.layers.attention.fla.chunk") - except ImportError as exc: - return None, f"cannot import sglang.srt.layers.attention.fla.chunk: {exc}" - if not hasattr(module, "chunk_gated_delta_rule"): - return None, "sglang.srt.layers.attention.fla.chunk has no chunk_gated_delta_rule" - return module.chunk_gated_delta_rule, "sglang.srt.layers.attention.fla.chunk.chunk_gated_delta_rule" - - -def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: - if tp_size not in (1, 2, 4, 8): - raise ValueError(f"tp_size must be one of 1, 2, 4, 8, got {tp_size}") - return Qwen35LinearAttentionConfig( - hidden_size=CONFIG.hidden_size // tp_size, - conv_kernel_size=CONFIG.conv_kernel_size, - num_k_heads=CONFIG.num_k_heads // tp_size, - num_v_heads=CONFIG.num_v_heads // tp_size, - head_k_dim=CONFIG.head_k_dim, - head_v_dim=CONFIG.head_v_dim, - qkv_dtype=CONFIG.qkv_dtype, - state_dtype=CONFIG.state_dtype, - ) - - -def make_inputs(batch: int, seq_len: int, *, device: torch.device, seed: int, config: Qwen35LinearAttentionConfig): - torch.manual_seed(seed) - q = torch.randn(batch, seq_len, config.num_v_heads, config.head_k_dim, device=device, dtype=config.qkv_dtype) - k = torch.randn_like(q) - v = torch.randn_like(q) - a = torch.randn(batch, seq_len, config.num_v_heads, device=device, dtype=config.qkv_dtype) - b = torch.randn(batch, seq_len, config.num_v_heads, device=device, dtype=config.qkv_dtype) - beta = torch.sigmoid(b.float()).to(dtype=config.qkv_dtype) - A_log = -torch.rand(config.num_v_heads, device=device, dtype=torch.float32) - dt_bias = torch.randn(config.num_v_heads, device=device, dtype=torch.float32) * 0.1 - log_gate = (-torch.exp(A_log).view(1, 1, -1) * torch.nn.functional.softplus(a.float() + dt_bias.view(1, 1, -1))).to( - dtype=config.qkv_dtype - ) - initial_state = torch.randn( - batch, - config.num_v_heads, - config.head_k_dim, - config.head_v_dim, - device=device, - dtype=torch.float32, - ) * 0.01 - mixed_qkv_conv = torch.randn(batch * seq_len, config.conv_dim, device=device, dtype=config.qkv_dtype) - a_flat = a.reshape(batch * seq_len, config.num_v_heads).contiguous() - b_flat = b.reshape(batch * seq_len, config.num_v_heads).contiguous() - return q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, mixed_qkv_conv, a_flat, b_flat - - -def run_cula_chunk_qk(q, k, out): - if cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): - raise RuntimeError("cula.cudac.qwen35_chunk_qk_prefill_sm90 is not available. Rebuild the CUDA extension.") - cula_cuda.qwen35_chunk_qk_prefill_sm90(q.contiguous(), k.contiguous(), out) - return out - - -def run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state): - return qwen35_scalar_kda_prefill( - q, - k, - v, - a, - b, - A_log, - dt_bias, - initial_state=initial_state, - backend="cudac", - ) - - -def run_cula_fused(q, k, v, a, b, A_log, dt_bias, initial_state): - return qwen35_fused_kda_prefill( - q, - k, - v, - a, - b, - A_log, - dt_bias, - initial_state=initial_state, - ) - - -def prepare_cula_fused_core_inputs(q, k, a, b, A_log, dt_bias, initial_state): - B, T, HV, K = q.shape - q_norm = F.normalize(q.float(), dim=-1).to(q.dtype).contiguous() - k_norm = F.normalize(k.float(), dim=-1).to(k.dtype).contiguous() - log_gate_scalar = -torch.exp(A_log.float()).view(1, 1, HV, 1) * F.softplus( - a.float().unsqueeze(-1) + dt_bias.float().view(1, 1, HV, 1) - ) - log_gate = log_gate_scalar.expand(B, T, HV, K).contiguous() - chunks = [] - for chunk_start in range(0, T, 64): - chunks.append(log_gate[:, chunk_start : chunk_start + 64].cumsum(dim=1) * RCP_LN2) - log_gate_cumsum = torch.cat(chunks, dim=1).contiguous() - beta = torch.sigmoid(b.float()).contiguous() - initial_state_vk = initial_state.float().transpose(-1, -2).contiguous() - return q_norm, k_norm, log_gate_cumsum, beta, initial_state_vk - - -def run_cula_fused_core(q_norm, k_norm, v, log_gate_cumsum, beta, initial_state_vk, config: Qwen35LinearAttentionConfig): - fused_kda_prefill = get_kda_fused_fwd(q_norm.device) - return fused_kda_prefill( - q=q_norm, - k=k_norm, - v=v.contiguous(), - g=log_gate_cumsum, - beta=beta, - scale=config.head_k_dim**-0.5, - initial_state=initial_state_vk, - output_final_state=True, - use_qk_l2norm_in_kernel=False, - use_gate_in_kernel=False, - safe_gate=False, - g_is_cumsum=True, - ) - - -def run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices, config: Qwen35LinearAttentionConfig): - # SGLang/FLA GDR chunk kernels use [N, H, V, K] state layout. cuLA's - # Qwen3.5 wrapper uses [N, H, K, V], so pass the transposed view here. - initial_state_vk = initial_state.transpose(-1, -2).contiguous() - kwargs = dict( - q=q, - k=k, - v=v, - g=log_gate, - beta=beta, - initial_state=initial_state_vk, - initial_state_indices=initial_state_indices, - output_final_state=True, - scale=config.head_k_dim**-0.5, - use_qk_l2norm_in_kernel=True, - head_first=False, - ) - try: - sig = inspect.signature(chunk_gdr) - kwargs = {key: value for key, value in kwargs.items() if key in sig.parameters} - except (TypeError, ValueError): - pass - return chunk_gdr(**kwargs) - - -def _normalize_chunk_result(result): - if isinstance(result, tuple): - out = result[0] - state = result[-1] if len(result) >= 2 else None - return out, state - return result, None - - -def _state_to_cula_layout(state: torch.Tensor | None) -> torch.Tensor | None: - if state is None: - return None - return state.transpose(-1, -2).contiguous() - - -def print_header(device: torch.device, args: argparse.Namespace, baseline_sources: dict[str, str]) -> None: - config = local_config_from_tp_size(args.tp_size) - print("Qwen3.5 prefill benchmark") - print(f" device: {torch.cuda.get_device_name(device)}") - print(f" dtype: {config.qkv_dtype}") - print(f" batch: {args.batch}") - print( - f" tp/local config: tp={args.tp_size} local_k_heads={config.num_k_heads} " - f"local_v_heads={config.num_v_heads} conv_dim={config.conv_dim}" - ) - print(f" seq lens: {args.seq_lens}") - print(f" warmup/rep: {args.warmup}/{args.rep}") - print(f" baselines: {baseline_sources or 'disabled/unavailable'}") - if args.cula_mode == "qk": - print(" cula: qwen35_chunk_qk_prefill_sm90 QK subpath only; baselines are full Triton GDR chunk kernels") - elif args.cula_mode == "scalar": - print(" cula: qwen35_scalar_kda_prefill full recurrence fallback") - elif args.cula_mode == "fused": - print(" cula: qwen35_fused_kda_prefill full recurrence via fused KDA CuTe core") - elif args.cula_mode == "fused-core": - print(" cula: fused KDA CuTe core only; Qwen gate/l2norm/cumsum/state prep is outside timing") - print() - cula_col = f"cula_{args.cula_mode}_ms" - print( - f"{'baseline':>8} {'B':>3} {'T':>7} {'layout_ms':>11} {cula_col:>13} {'cula_total':>11} " - f"{'base_ms':>11} {'base/cula':>10} {'rel_rms':>11} {'rel_max':>11}" - ) - print("-" * 113) - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--batch", type=int, default=1) - parser.add_argument("--seq-lens", type=int, nargs="+", default=[128, 256, 512, 1024]) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--rep", type=int, default=30) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--tp-size", type=int, choices=[1, 2, 4, 8], default=1) - parser.add_argument("--baseline", choices=["none", "fla", "sgl", "all"], default="sgl") - parser.add_argument("--sglang-path", type=pathlib.Path, default=None) - parser.add_argument( - "--cula-mode", - choices=["qk", "scalar", "fused", "fused-core"], - default="qk", - help="cuLA path to benchmark: qk is QK subpath, scalar is old full fallback, fused is wrapper, fused-core is kernel only.", - ) - parser.add_argument("--skip-accuracy", action="store_true") - parser.add_argument( - "--max-qk-elements", - type=int, - default=512 * 1024 * 1024, - help="Skip cuLA QK timings when B*local_HV*T*T exceeds this element count.", - ) - args = parser.parse_args() - config = local_config_from_tp_size(args.tp_size) - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for this benchmark.") - device = torch.device("cuda") - - baselines: dict[str, Callable] = {} - baseline_sources: dict[str, str] = {} - if args.baseline in ("fla", "all"): - fla_chunk_gdr, fla_source_or_error = resolve_fla_chunk_gdr() - if fla_chunk_gdr is None: - print(f"Skipping FLA baseline: {fla_source_or_error}") - else: - baselines["fla"] = fla_chunk_gdr - baseline_sources["fla"] = fla_source_or_error - if args.baseline in ("sgl", "all"): - sgl_chunk_gdr, sgl_source_or_error = resolve_sgl_chunk_gdr(args.sglang_path) - if sgl_chunk_gdr is None: - print(f"Skipping SGLang baseline: {sgl_source_or_error}") - else: - baselines["sgl"] = sgl_chunk_gdr - baseline_sources["sgl"] = sgl_source_or_error - - print_header(device, args, baseline_sources) - - for seq_len in args.seq_lens: - q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, mixed_qkv_conv, a_flat, b_flat = make_inputs( - args.batch, - seq_len, - device=device, - seed=args.seed, - config=config, - ) - initial_state_indices = torch.arange(args.batch, device=device, dtype=torch.int32) - - def layout_fn(): - return qwen35_layout_prefill(mixed_qkv_conv, a_flat, b_flat, backend="cudac") - - qk_elements = args.batch * config.num_v_heads * seq_len * seq_len - qk_out = None - if args.cula_mode == "qk" and qk_elements <= args.max_qk_elements: - qk_out = torch.empty( - args.batch, - config.num_v_heads, - seq_len, - seq_len, - device=device, - dtype=torch.float32, - ) - fused_core_inputs = None - if args.cula_mode == "fused-core": - fused_core_inputs = prepare_cula_fused_core_inputs(q, k, a, b, A_log, dt_bias, initial_state) - - def cula_fn(): - if args.cula_mode == "scalar": - return run_cula_scalar(q, k, v, a, b, A_log, dt_bias, initial_state) - if args.cula_mode == "fused": - return run_cula_fused(q, k, v, a, b, A_log, dt_bias, initial_state) - if args.cula_mode == "fused-core": - return run_cula_fused_core(*fused_core_inputs[:2], v, *fused_core_inputs[2:], config) - if qk_out is None: - raise RuntimeError( - f"Skipping cuLA QK: B*H*T*T={qk_elements} exceeds --max-qk-elements={args.max_qk_elements}" - ) - return run_cula_chunk_qk(q, k, qk_out) - - layout_ms = benchmark_cuda_fn(layout_fn, warmup=args.warmup, rep=args.rep) - cula_ms = ( - float("nan") - if args.cula_mode == "qk" and qk_out is None - else benchmark_cuda_fn(cula_fn, warmup=args.warmup, rep=args.rep) - ) - cula_total_ms = layout_ms + cula_ms if not torch.isnan(torch.tensor(cula_ms)) else float("nan") - - rel_rms = float("nan") - rel_max = float("nan") - state_cula = None - if not args.skip_accuracy: - out_cula = cula_fn() - if args.cula_mode == "qk": - qk_ref = torch.einsum("bthd,bshd->bhts", q.float(), k.float()) - torch.cuda.synchronize() - rel_rms, rel_max, _ = error_stats(qk_ref, out_cula) - del qk_ref - else: - out_cula, state_cula = out_cula - torch.cuda.synchronize() - - if not baselines: - print( - f"{'none':>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_ms:13.4f} {cula_total_ms:11.4f} " - f"{float('nan'):11.4f} {float('nan'):10.3f} {rel_rms:11.3e} {rel_max:11.3e}" - ) - - for baseline_name, chunk_gdr in baselines.items(): - def baseline_fn(): - return run_chunk_gdr(chunk_gdr, q, k, v, log_gate, beta, initial_state, initial_state_indices, config) - - row_rel_rms = rel_rms - row_rel_max = rel_max - if args.cula_mode in ("scalar", "fused", "fused-core") and not args.skip_accuracy: - if state_cula is None: - out_cula, state_cula = cula_fn() - out_base, state_base = _normalize_chunk_result(baseline_fn()) - state_base = _state_to_cula_layout(state_base) - torch.cuda.synchronize() - row_rel_rms, row_rel_max, _ = error_stats(out_base, out_cula) - if state_base is not None and tuple(state_base.shape) == tuple(state_cula.shape): - rel_rms_s, rel_max_s, _ = error_stats(state_base, state_cula) - row_rel_rms = max(row_rel_rms, rel_rms_s) - row_rel_max = max(row_rel_max, rel_max_s) - - base_ms = benchmark_cuda_fn(baseline_fn, warmup=args.warmup, rep=args.rep) - speedup = base_ms / cula_ms if cula_ms > 0 else float("nan") - print( - f"{baseline_name:>8} {args.batch:3d} {seq_len:7d} {layout_ms:11.4f} {cula_ms:13.4f} {cula_total_ms:11.4f} " - f"{base_ms:11.4f} {speedup:10.3f} {row_rel_rms:11.3e} {row_rel_max:11.3e}" - ) - - del q, k, v, a, b, beta, log_gate, A_log, dt_bias, initial_state, initial_state_indices, mixed_qkv_conv, a_flat, b_flat, qk_out, fused_core_inputs - torch.cuda.empty_cache() - - -if __name__ == "__main__": - main() diff --git a/cula/kda/blackwell_fused_fwd.py b/cula/kda/blackwell_fused_fwd.py index 5ac821ba..595f7aad 100644 --- a/cula/kda/blackwell_fused_fwd.py +++ b/cula/kda/blackwell_fused_fwd.py @@ -115,11 +115,13 @@ def forward( chunk_indices: torch.IntTensor | None = None, ): chunk_size = 64 - assert q.shape[-2] == v.shape[-2] == k.shape[-2], "Number of heads must be the same for q, k, v." + assert q.shape == k.shape, "q and k must have the same shape." global compiled_kernel_cache B, S, H, D = q.shape + HV = v.shape[-2] + assert HV % H == 0, f"HV ({HV}) must be a multiple of H ({H})." is_varlen = cu_seqlens is not None if is_varlen: assert B == 1, "For varlen, batch size must be 1. Flatten variable-length inputs first." @@ -168,7 +170,7 @@ def forward( g_cute = from_dlpack(g.detach()) beta_cute = from_dlpack(beta.detach()) - o = torch.empty_like(q) + o = torch.empty_like(v) o_cute = from_dlpack(o.detach()) stream = cutlass_torch.default_stream() @@ -227,13 +229,13 @@ def forward( initial_state_cute = _dummy_cache[q.device]["state_cute"] if output_final_state: - final_state_f32 = torch.zeros(num_seqs, H, D, D, dtype=torch.float32, device=q.device) + final_state_f32 = torch.zeros(num_seqs, HV, D, D, dtype=torch.float32, device=q.device) final_state_cute = from_dlpack(final_state_f32.detach()) else: final_state_f32 = None final_state_cute = _dummy_cache[q.device]["state_cute"] - problem_size = (num_seqs, S, H, D) + problem_size = (num_seqs, S, H, HV, D) if cache_key in compiled_kernel_cache: compiled_kernel = compiled_kernel_cache[cache_key] diff --git a/cula/ops/kda_fully_fused_sm100_wip.py b/cula/ops/kda_fully_fused_sm100_wip.py index ec3d5fe9..7f9407fa 100644 --- a/cula/ops/kda_fully_fused_sm100_wip.py +++ b/cula/ops/kda_fully_fused_sm100_wip.py @@ -56,6 +56,7 @@ import argparse import time +from types import SimpleNamespace import cuda.bindings.driver as cuda import cutlass @@ -68,6 +69,14 @@ from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack from cutlass.cute.typing import Int32, Int64 + +# CUTLASS DSL 4.3+ changed fence_proxy() from exported enum arguments to +# string literals. Keep the existing call sites compatible with both APIs. +if not hasattr(cute.arch, "ProxyKind"): + cute.arch.ProxyKind = SimpleNamespace(async_shared="async.shared") +if not hasattr(cute.arch, "SharedSpace"): + cute.arch.SharedSpace = SimpleNamespace(shared_cta="cta") + try: from fla.modules.l2norm import l2norm_fwd except ImportError: @@ -373,7 +382,7 @@ def __call__( final_state_iter: cute.Pointer, # Final state [B, H, D, D], float32 or nullptr cu_seqlens_iter: cute.Pointer, # Cumulative seq lengths [num_seqs+1], int32 (varlen) workspace_iter: cute.Pointer, # Workspace buffer for TMA descriptor modification - problem_size: tuple[Int32, Int32, Int32, Int32], # (B/num_seqs, S/total_tokens, H, D) + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], # (B/num_seqs, S/total_tokens, H, HV, D) stream: cuda.CUstream, options=None, # compile options ): @@ -393,11 +402,11 @@ def __call__( final_state_iter: Final state [N, H, D, D] or nullptr cu_seqlens_iter: Cumulative seq lengths [num_seqs+1], int32 (varlen only) workspace_iter: Workspace buffer for TMA descriptor modification (varlen tail tiles) - problem_size: (N, S, H, D) where N=B or num_seqs, S=seq_len or total_tokens + problem_size: (N, S, H, HV, D) where N=B or num_seqs, S=seq_len or total_tokens stream: CUDA stream options: compile options for the kernel """ - B, S, H, D = problem_size + B, S, H, HV, D = problem_size # Setup attributes self._setup_attributes() @@ -433,36 +442,36 @@ def __call__( kt = cute.make_tensor(k_iter, kt_layout) # v v_layout = cute.make_layout( - (D, S, (H, data_B)), - stride=(1, D * H, (D, D * H * S)), + (D, S, (HV, data_B)), + stride=(1, D * HV, (D, D * HV * S)), ) v = cute.make_tensor(v_iter, v_layout) # g (gate) - NEW for KDA, same layout as Q/K g_layout = cute.make_layout( - (S, D, (H, data_B)), - stride=(D * H, 1, (D, D * H * S)), + (S, D, (HV, data_B)), + stride=(D * HV, 1, (D, D * HV * S)), ) g = cute.make_tensor(g_iter, g_layout) # beta - NEW for KDA, shape (B, S, H) or (1, total_tokens, H) for varlen beta_layout = cute.make_layout( - (S, (H, data_B)), - stride=(H, (1, H * S)), + (S, (HV, data_B)), + stride=(HV, (1, HV * S)), ) beta = cute.make_tensor(beta_iter, beta_layout) o_layout = cute.make_layout( - (D, S, (H, data_B)), - stride=(1, D * H, (D, D * H * S)), + (D, S, (HV, data_B)), + stride=(1, D * HV, (D, D * HV * S)), ) o = cute.make_tensor(o_iter, o_layout) # Initial state / final state: [N, H, D, D] stored as row-major # N = B (non-varlen) or num_seqs (varlen). Always uses B from problem_size. fstate_layout = cute.make_layout( - (D, D, (H, B)), - stride=(1, D, (D * D, D * D * H)), + (D, D, (HV, B)), + stride=(1, D, (D * D, D * D * HV)), ) initial_state = cute.make_tensor(initial_state_iter, fstate_layout) final_state = cute.make_tensor(final_state_iter, fstate_layout) @@ -837,7 +846,7 @@ class SharedStorage: self.shared_storage = SharedStorage if cutlass.const_expr(self.is_varlen): - self.grid = (1, H, B) + self.grid = (1, HV, B) # TensorMapManager for TMA descriptor modification in varlen tail tiles self._tensormap_mgr = utils.TensorMapManager(utils.TensorMapUpdateMode.GMEM, 128) else: @@ -931,7 +940,7 @@ def kernel( cu_seqlens: cute.Tensor, # int32 tensor for varlen o_gmem: cute.Tensor, # raw GMEM output tensor (D, S, (H, data_B)) for tail tile handling workspace_iter: cute.Pointer, # workspace buffer for TMA descriptor modification - problem_size: tuple[Int32, Int32, Int32, Int32], # (B, S, H, D) + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], # (B, S, H, HV, D) ): """ KDA Kernel - Step 1: Gate processing @@ -1381,7 +1390,8 @@ def kernel( ) (_, hidx, bidx) = cute.arch.block_idx() - B, S, H, D = problem_size + B, S, H, HV, D = problem_size + qk_hidx = hidx // (HV // H) C = self.chunk_size # Varlen: compute per-CTA sequence boundary and domain offsets @@ -1546,6 +1556,7 @@ def kernel( operand_mode="A", debug_name="Q", batch_idx=data_bidx, + head_idx=qk_hidx, ) tKsK, tKgK = self.tma_partition_for_mma_operand( @@ -1557,6 +1568,7 @@ def kernel( operand_mode="B", debug_name="K", batch_idx=data_bidx, + head_idx=qk_hidx, ) tVsV, tVgV = self.tma_partition_for_mma_operand( @@ -4098,8 +4110,8 @@ def index_transform(index_q, index_k): o_tail = cute.make_tensor( o_gmem.iterator, cute.make_layout( - (D, new_S, (H, Int32(1))), - stride=(Int32(1), D * H, (D, D * H * S)), + (D, new_S, (HV, Int32(1))), + stride=(Int32(1), D * HV, (D, D * HV * S)), ), ) # Initialize: copy original TMA descriptor to workspace @@ -4157,8 +4169,8 @@ def index_transform(index_q, index_k): if cutlass.const_expr(self.is_varlen): beta_v = cute.domain_offset((tok_offset, (0, 0)), beta) beta_chunk = beta_v[(None, (hidx, data_bidx))] - beta_chunk_layout = cute.make_layout((C, 1), stride=(H, 0)) - beta_chunk = cute.make_tensor(beta_chunk.iterator + s_idx * H, layout=beta_chunk_layout) + beta_chunk_layout = cute.make_layout((C, 1), stride=(HV, 0)) + beta_chunk = cute.make_tensor(beta_chunk.iterator + s_idx * HV, layout=beta_chunk_layout) if cutlass.const_expr(PRINT_DEBUG): print(f"sBeta: {sBeta}") @@ -5513,10 +5525,13 @@ def local_tile_partition_for_mma_operand( debug_name=None, no_cta_coord=False, batch_idx=None, + head_idx=None, ): _, hidx, bidx = cute.arch.block_idx() if batch_idx is not None: bidx = batch_idx + if cutlass.const_expr(head_idx is not None): + hidx = head_idx # Local_tile partition global tensors # x: (0,0,0,0) o (M,K,(H,B)):(1@1,1@0,(1@2,1@3)) # (MMATile_M, MMATile_K, TILES_M, TILES_K, (H, B)) @@ -5570,6 +5585,7 @@ def tma_partition_for_mma_operand( operand_mode, debug_name=None, batch_idx=None, + head_idx=None, ): tCgX = self.local_tile_partition_for_mma_operand( tensor_x=tma_tensor_x, @@ -5578,6 +5594,7 @@ def tma_partition_for_mma_operand( operand_mode=operand_mode, debug_name=debug_name, batch_idx=batch_idx, + head_idx=head_idx, ) # Partition shared tensor with regard to TMA # ((ATOM_V, REST_V), INPUT_STAGE) From aa3d9330669d9d33d81c0ff6596818f8921b46fd Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sun, 2 Aug 2026 15:38:46 +0000 Subject: [PATCH 20/35] perf(qwen35): optimize native GVA fused prefill --- cula/kda/blackwell_fused_fwd.py | 76 +- cula/ops/kda_fully_fused_sm100_wip.py | 1330 ++++++++++++++++++------- cula/ops/qwen35_fused_kda_prefill.py | 54 +- 3 files changed, 1062 insertions(+), 398 deletions(-) diff --git a/cula/kda/blackwell_fused_fwd.py b/cula/kda/blackwell_fused_fwd.py index 595f7aad..f2381021 100644 --- a/cula/kda/blackwell_fused_fwd.py +++ b/cula/kda/blackwell_fused_fwd.py @@ -14,6 +14,7 @@ import pathlib import sys +import types import warnings import torch @@ -33,6 +34,8 @@ from fla.ops.utils import chunk_local_cumsum from fla.ops.utils.constant import RCP_LN2 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + from cula.ops.l2norm_triton import l2norm_fwd, l2norm_qk_fwd + except ImportError: RCP_LN2 = 1.4426950408889634 @@ -49,6 +52,11 @@ def l2norm_fwd(x: torch.Tensor): rstd = torch.rsqrt(x.float().square().sum(dim=-1, keepdim=True).clamp_min(1.0e-12)) return (x.float() * rstd).to(x.dtype), rstd + def l2norm_qk_fwd(q: torch.Tensor, k: torch.Tensor): + q_out, q_rstd = l2norm_fwd(q) + k_out, k_rstd = l2norm_fwd(k) + return q_out, k_out, q_rstd, k_rstd + def kda_gate_fwd(*args, **kwargs): raise ImportError("fla is required for use_gate_in_kernel=True in blackwell_fused_fwd") @@ -111,6 +119,7 @@ def forward( safe_gate: bool = False, lower_bound: float | None = None, g_is_cumsum: bool = False, + scalar_gate: bool = False, cu_seqlens: torch.IntTensor | None = None, chunk_indices: torch.IntTensor | None = None, ): @@ -128,6 +137,16 @@ def forward( num_seqs = cu_seqlens.shape[0] - 1 else: num_seqs = B + split_value_tiles = ( + scalar_gate + and safe_gate + and not is_varlen + and B == 1 + and H == 16 + and HV == 48 + and D == 128 + and v.shape[-1] == 128 + ) g_org = None if use_gate_in_kernel: @@ -155,14 +174,22 @@ def forward( A_log=A_log, dt_bias=dt_bias, ) - if not g_is_cumsum and not (safe_gate and use_gate_in_kernel): + fuse_scalar_cumsum = ( + scalar_gate + and not g_is_cumsum + and not use_gate_in_kernel + ) + if ( + not g_is_cumsum + and not (safe_gate and use_gate_in_kernel) + and not fuse_scalar_cumsum + ): g = chunk_local_cumsum( g=g, chunk_size=chunk_size, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices ) q_rstd, k_rstd = None, None if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) + q, k, q_rstd, k_rstd = l2norm_qk_fwd(q, k) q_cute = from_dlpack(q.detach()) k_cute = from_dlpack(k.detach()) @@ -176,7 +203,24 @@ def forward( stream = cutlass_torch.default_stream() has_initial_state = initial_state is not None - cache_key = (has_initial_state, output_final_state, safe_gate, is_varlen, scale, chunk_size, D, USE_FAST_MATH) + # H/HV affect the compiled tensor layouts and GVA head mapping. Without + # them, a process exercising multiple Qwen model/TP shapes can reuse a + # kernel compiled for a different head configuration. + cache_key = ( + has_initial_state, + output_final_state, + safe_gate, + scalar_gate, + fuse_scalar_cumsum, + split_value_tiles, + is_varlen, + scale, + chunk_size, + H, + HV, + D, + USE_FAST_MATH, + ) if is_varlen: cu_seqlens_i32 = cu_seqlens.to(torch.int32).contiguous() @@ -229,7 +273,9 @@ def forward( initial_state_cute = _dummy_cache[q.device]["state_cute"] if output_final_state: - final_state_f32 = torch.zeros(num_seqs, HV, D, D, dtype=torch.float32, device=q.device) + # The CuTe kernel overwrites every final-state element. Avoid a + # redundant ~3 MiB memset for Qwen3.5-27B (HV=48) on every call. + final_state_f32 = torch.empty(num_seqs, HV, D, D, dtype=torch.float32, device=q.device) final_state_cute = from_dlpack(final_state_f32.detach()) else: final_state_f32 = None @@ -247,6 +293,9 @@ def forward( io_dtype=cutlass.BFloat16, scale=scale, safe_gate=safe_gate, + scalar_gate=scalar_gate, + fuse_scalar_cumsum=fuse_scalar_cumsum, + split_value_tiles=split_value_tiles, has_initial_state=has_initial_state, output_final_state=output_final_state, is_varlen=is_varlen, @@ -317,6 +366,7 @@ def flash_kda_prefill( safe_gate: bool = False, lower_bound: float | None = None, g_is_cumsum: bool = False, + scalar_gate: bool = False, cu_seqlens: torch.IntTensor | None = None, chunk_indices: torch.IntTensor | None = None, **kwargs, @@ -356,12 +406,13 @@ def flash_kda_prefill( assert HV % H == 0, ( f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by num_qk_heads (H={H}), but got HV % H = {HV % H}" ) - assert g.shape == (B, T, HV, K), f"g must have shape [B, T, HV, K]={[B, T, HV, K]}, got {list(g.shape)}" + expected_g_shape = (B, T, HV) if scalar_gate else (B, T, HV, K) + assert g.shape == expected_g_shape, f"g must have shape {expected_g_shape}, got {list(g.shape)}" assert beta.shape == (B, T, HV), f"beta must have shape [B, T, HV]={[B, T, HV]}, got {list(beta.shape)}" if scale is None: scale = k.shape[-1] ** -0.5 - o, final_state = ChunkKDAFunction.apply( + forward_args = ( q, k, v, @@ -377,7 +428,18 @@ def flash_kda_prefill( safe_gate, lower_bound, g_is_cumsum, + scalar_gate, cu_seqlens, chunk_indices, ) + if torch.is_grad_enabled() and any( + tensor is not None and tensor.requires_grad + for tensor in (q, k, v, g, beta, initial_state) + ): + o, final_state = ChunkKDAFunction.apply(*forward_args) + else: + # The op has no backward implementation. SGLang always reaches this + # inference branch, so avoid the measurable autograd.Function.apply + # dispatch cost while retaining AMP/input-guard behavior. + o, final_state = ChunkKDAFunction.forward(types.SimpleNamespace(), *forward_args) return o, final_state diff --git a/cula/ops/kda_fully_fused_sm100_wip.py b/cula/ops/kda_fully_fused_sm100_wip.py index 7f9407fa..498ea40d 100644 --- a/cula/ops/kda_fully_fused_sm100_wip.py +++ b/cula/ops/kda_fully_fused_sm100_wip.py @@ -104,6 +104,7 @@ class Constant: D = 128 # head dim HALF_D = 64 # half head dim for partitioned S2R SCALE = float(D) ** -0.5 + RCP_LN2 = 1.4426950408889634 BK_SC = 64 # tile size in subchunk MMA @@ -140,19 +141,32 @@ def __init__( io_dtype: type[cutlass.Numeric] = cutlass.BFloat16, scale: cutlass.Float32 = 1.0, safe_gate: bool = False, + scalar_gate: bool = False, + fuse_scalar_cumsum: bool = False, + split_value_tiles: bool = False, has_initial_state: bool = False, output_final_state: bool = False, is_varlen: bool = False, use_fast_math: bool = True, - # num_regs_cuda: int = 248, num_regs_cuda: int = 224, num_regs_subchunk: int = 192, - num_regs_others: int = 64, # Optimized: best config from comprehensive sweep + num_regs_others: int = 80, ): assert_blackwell() # make scale a constant self.scale = scale self.safe_gate = safe_gate + self.scalar_gate = scalar_gate + self.fuse_scalar_cumsum = fuse_scalar_cumsum + self.split_value_tiles = split_value_tiles + if scalar_gate and not safe_gate: + raise ValueError("native scalar gate currently requires safe_gate=True") + if fuse_scalar_cumsum and not scalar_gate: + raise ValueError("fuse_scalar_cumsum requires scalar_gate=True") + if split_value_tiles and (not scalar_gate or not safe_gate or is_varlen): + raise ValueError( + "split_value_tiles currently requires non-varlen safe scalar gate" + ) self.has_initial_state = has_initial_state self.output_final_state = output_final_state self.is_varlen = is_varlen @@ -183,6 +197,9 @@ def __init__( # K: (64, 128) # V: (64, 128) C, D = (Constant.C, Constant.D) + self.value_tile_size = D // 2 if split_value_tiles else D + self.num_value_tiles = D // self.value_tile_size + DV_TILE = self.value_tile_size HALF_D = Constant.HALF_D # (C, C, D) self.qk_mma_tiler = (C, C, D) # (M, N, K) @@ -190,15 +207,15 @@ def __init__( self.qk_mma_tiler_half = (C, C, HALF_D) # (M, N, K/2) self.kk_mma_tiler = (C, C, D) # (M, N, K) # (D, C, C) - self.vp_mma_tiler = (D, C, C) # (M, N, K) - self.mv_mma_tiler = (D, C, C) # (M, N, K) + self.vp_mma_tiler = (DV_TILE, C, C) # (M, N, K) + self.mv_mma_tiler = (DV_TILE, C, C) # (M, N, K) # (D, D, C) - self.kv_mma_tiler = (D, D, C) # (M, N, K) + self.kv_mma_tiler = (DV_TILE, D, C) # (M, N, K) # (D, C, D) # State as operand A since it's in TMEM # Q now as operand B - self.sq_mma_tiler = (D, C, D) # (M, N, K) - self.ks_mma_tiler = (D, C, D) # (M, N, K) + self.sq_mma_tiler = (DV_TILE, C, D) # (M, N, K) + self.ks_mma_tiler = (DV_TILE, C, D) # (M, N, K) # subchunk MMA SC, BK_SC = (Constant.SC, Constant.BK_SC) @@ -340,8 +357,8 @@ def _setup_attributes(self): self.k_stage = 2 self.v_stage = 1 self.o_stage = 2 - self.g_stage = 2 # Single stage for g (CUDA warp processes immediately) - self.beta_stage = 1 # TODO: two stage ? + self.g_stage = 2 + self.beta_stage = 2 self.q_k_scaled_stage = 1 # only single stage here due to smem limitation self.epi_stage = 2 @@ -362,7 +379,7 @@ def _compute_grid( # cute.ceil_div(o_shape[0], chunk_size), # For Loop to tile over chunk size, # TODO: varlen will make parallelism good enough - 1, + self.num_value_tiles, # H cute.size(o_shape[2][0]), # B @@ -435,11 +452,6 @@ def __call__( stride=(D * H, 1, (D, D * H * S)), ) k = cute.make_tensor(k_iter, k_layout) - kt_layout = cute.make_layout( - (D, S, (H, data_B)), - stride=(1, D * H, (D, D * H * S)), - ) - kt = cute.make_tensor(k_iter, kt_layout) # v v_layout = cute.make_layout( (D, S, (HV, data_B)), @@ -447,11 +459,19 @@ def __call__( ) v = cute.make_tensor(v_iter, v_layout) - # g (gate) - NEW for KDA, same layout as Q/K - g_layout = cute.make_layout( - (S, D, (HV, data_B)), - stride=(D * HV, 1, (D, D * HV * S)), - ) + # GDN uses one scalar gate per token/value-head. Keep a singleton + # second mode so the scalar specialization can use a regular 2-D TMA + # tile while preserving the generic vector-gate layout unchanged. + if cutlass.const_expr(self.scalar_gate): + g_layout = cute.make_layout( + (S, 1, (HV, data_B)), + stride=(HV, 0, (1, HV * S)), + ) + else: + g_layout = cute.make_layout( + (S, D, (HV, data_B)), + stride=(D * HV, 1, (D, D * HV * S)), + ) g = cute.make_tensor(g_iter, g_layout) # beta - NEW for KDA, shape (B, S, H) or (1, total_tokens, H) for varlen @@ -491,7 +511,12 @@ def __call__( self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() - self.g_major_mode = utils.LayoutEnum.from_tensor(g).mma_major_mode() # NEW for KDA + # Scalar G is not an MMA operand. The value is only needed for the + # generic vector-gate TMA construction below. + if cutlass.const_expr(self.scalar_gate): + self.g_major_mode = self.q_major_mode + else: + self.g_major_mode = utils.LayoutEnum.from_tensor(g).mma_major_mode() self.k_major_mode_kv = tcgen05.OperandMajorMode.MN # For V^T*K, S dimension coalesced # TMEM register output results as (D, C) self.o_layout = utils.LayoutEnum.from_tensor(o) @@ -638,16 +663,30 @@ def __call__( self.k_dtype, self.k_stage, ) - # G (gate) - NEW for KDA - # Use same layout as Q since g has same shape and memory layout as Q - # This ensures TMA compatibility - # ((MMA_ATOM_M, MMA_ATOM_K), MMA_M, MMA_K, STAGES) - g_smem_layout_staged = sm100_utils.make_smem_layout_a( - qk_tiled_mma, - self.qk_mma_tiler, - self.g_dtype, - self.g_stage, - ) + # The scalar specialization stages only C values. Its storage is + # subsequently reused as BF16 K*exp(-g), so SharedStorage below keeps + # enough physical bytes for that tensor without retaining the 64 KiB + # FP32 vector-gate allocation. + if cutlass.const_expr(self.scalar_gate): + scalar_g_stage_elements = Constant.C * ( + 4 + Constant.C // Constant.SC + ) + (Constant.C // Constant.SC) ** 2 + g_smem_layout_staged = cute.make_composed_layout( + cute.make_swizzle(0, 4, 3), + 0, + cute.make_layout( + (Constant.C, 1, self.g_stage), + stride=(1, Constant.C, scalar_g_stage_elements), + ), + ) + else: + # Generic vector gate: same MMA-compatible layout as Q. + g_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.g_dtype, + self.g_stage, + ) # V^T*P p_smem_layout_staged = sm100_utils.make_smem_layout_b( vp_tiled_mma, @@ -702,15 +741,6 @@ def __call__( qk_tiled_mma, cluster_layout_vmnk.shape, ) - kv_k_smem_layout = cute.select(kv_k_smem_layout_staged, mode=[0, 1, 2]) - tma_atom_kt, tma_tensor_kt = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, - kt, - kv_k_smem_layout, - self.kv_mma_tiler, - kv_tiled_mma, - cluster_layout_vmnk.shape, - ) # TMA load for V v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_A( @@ -721,17 +751,26 @@ def __call__( vp_tiled_mma, cluster_layout_vmnk.shape, ) - # TMA load for G (gate) - NEW for KDA - # Use same TMA atom as Q since g has same layout as Q - g_smem_layout = cute.select(g_smem_layout_staged, mode=[0, 1, 2]) - tma_atom_g, tma_tensor_g = cute.nvgpu.make_tiled_tma_atom_A( - tma_load_op, - g, - g_smem_layout, - self.qk_mma_tiler, - qk_tiled_mma, - cluster_layout_vmnk.shape, - ) + # TMA load for G. Scalar G uses a (C, 1) epilogue-style tile; vector + # G retains the original MMA-operand descriptor. + if cutlass.const_expr(self.scalar_gate): + g_smem_layout = cute.select(g_smem_layout_staged, mode=[0, 1]) + tma_atom_g, tma_tensor_g = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_op, + g, + g_smem_layout, + (Constant.C, 1), + ) + else: + g_smem_layout = cute.select(g_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_g, tma_tensor_g = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + g, + g_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + cluster_layout_vmnk.shape, + ) # NOTE: G's last row will be extracted from sG in CUDA warp after TMA load # No separate TMA needed for G last row - we extract it from the full G tile @@ -745,16 +784,26 @@ def __call__( q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) - g_copy_size = cute.size_in_bytes(self.g_dtype, g_smem_layout) # NEW for KDA + v_copy_size = cute.size_in_bytes(self.v_dtype, v_smem_layout) + g_copy_size = cute.size_in_bytes(self.g_dtype, g_smem_layout) self.tma_copy_q_bytes = q_copy_size self.tma_copy_k_bytes = k_copy_size - # self.tma_copy_v_bytes = v_copy_size - self.tma_copy_v_bytes = k_copy_size + self.tma_copy_v_bytes = v_copy_size self.tma_copy_g_bytes = g_copy_size # NEW for KDA beta_layout = cute.make_layout((Constant.C, self.beta_stage), stride=(1, Constant.C)) g_last_layout = cute.make_layout((Constant.D, self.g_stage), stride=(1, Constant.D)) + # Per stage: raw G[C], row factors[C], and four sets of column + # factors[4,C]. This is 1.5 KiB/stage, still far below vector G. + if cutlass.const_expr(self.scalar_gate): + scalar_g_stage_elements = Constant.C * ( + 4 + Constant.C // Constant.SC + ) + (Constant.C // Constant.SC) ** 2 + g_storage_elements = scalar_g_stage_elements * self.g_stage + else: + g_storage_elements = cute.cosize(g_smem_layout_staged) + @cute.struct class SharedStorage: # Pipeline barriers @@ -790,8 +839,6 @@ class SharedStorage: ks_mbar_ptr: cute.struct.MemRange[Int64, self.ks_stage * 2] # type: ignore o_inter_mbar_ptr: cute.struct.MemRange[Int64, 1 * 2] # type: ignore smem_o_mbar_ptr: cute.struct.MemRange[Int64, self.acc_stage * 2] # type: ignore - kv_decay_mbar_ptr: cute.struct.MemRange[Int64, 1 * 2] # type: ignore - kv_decay_mbar_ptr: cute.struct.MemRange[Int64, 1 * 2] # type: ignore # Tmem holding buffer tmem_holding_buf: Int32 # Smem tensors @@ -818,7 +865,7 @@ class SharedStorage: ] # G (gate) - NEW for KDA sG: cute.struct.Align[ - cute.struct.MemRange[self.g_dtype, cute.cosize(g_smem_layout_staged)], # type: ignore + cute.struct.MemRange[self.g_dtype, g_storage_elements], # type: ignore self.buffer_align_bytes, ] # Store QK @@ -867,12 +914,11 @@ class SharedStorage: tma_tensor_q, tma_atom_k, tma_tensor_k, - tma_atom_kt, - tma_tensor_kt, tma_atom_v, tma_tensor_v, tma_atom_g, # NEW for KDA tma_tensor_g, # NEW for KDA + g, tma_atom_o, tma_tensor_o, beta, # NEW for KDA @@ -915,12 +961,11 @@ def kernel( tma_tensor_q: cute.Tensor, tma_atom_k: cute.CopyAtom, tma_tensor_k: cute.Tensor, - tma_atom_kt: cute.CopyAtom, - tma_tensor_kt: cute.Tensor, tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, tma_atom_g: cute.CopyAtom, # NEW for KDA tma_tensor_g: cute.Tensor, # NEW for KDA + g: cute.Tensor, tma_atom_o: cute.CopyAtom, tma_tensor_o: cute.Tensor, beta: cute.Tensor, # NEW for KDA - shape (S, (H, B)) @@ -961,7 +1006,8 @@ def kernel( cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) - cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_g) # NEW for KDA + if cutlass.const_expr(not self.scalar_gate): + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_g) cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_o) # Allocate shared memory @@ -1036,16 +1082,27 @@ def kernel( consumer_group=make_thread_cooperative_group(32 * len(self.cuda_warp_ids)), barrier_storage=storage.end_v_mbar_ptr.data_ptr(), ).make_participants() - # G (gate/g_cumsum) - NEW for KDA - load_g_producer, load_g_consumer = pipeline.PipelineTmaAsync.create( - num_stages=self.g_stage, - producer_group=make_thread_cooperative_group(len([self.load_warp_id])), - consumer_group=make_thread_cooperative_group( - len([*self.cuda_warp_ids, *self.cuda_subchunk_warp_ids]) - ), # CUDA cores will consume - tx_count=self.tma_copy_g_bytes, - barrier_storage=storage.load_g_mbar_ptr.data_ptr(), - ).make_participants() + # Scalar G is copied by the load warp (two values per lane); vector G + # retains its TMA pipeline. + if cutlass.const_expr(self.scalar_gate): + load_g_producer, load_g_consumer = pipeline.PipelineAsync.create( + num_stages=self.g_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len([*self.cuda_warp_ids, *self.cuda_subchunk_warp_ids]) + ), + barrier_storage=storage.load_g_mbar_ptr.data_ptr(), + ).make_participants() + else: + load_g_producer, load_g_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.g_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group( + len([*self.cuda_warp_ids, *self.cuda_subchunk_warp_ids]) + ), + tx_count=self.tma_copy_g_bytes, + barrier_storage=storage.load_g_mbar_ptr.data_ptr(), + ).make_participants() load_beta_producer, load_beta_consumer = pipeline.PipelineAsync.create( num_stages=self.beta_stage, producer_group=make_thread_cooperative_group(self.threads_per_warp * len([self.load_beta_warp_id])), @@ -1153,14 +1210,6 @@ def kernel( consumer_group=make_thread_cooperative_group(self.threads_per_warp * len([self.epilogue_warp_id])), barrier_storage=storage.smem_o_mbar_ptr.data_ptr(), ).make_participants() - # T2R & R2T sync in S decay - kv_decay_producer, kv_decay_consumer = pipeline.PipelineAsyncUmma.create( - num_stages=1, - producer_group=make_thread_cooperative_group(self.threads_per_warp * len(self.cuda_warp_ids)), - consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), - barrier_storage=storage.kv_decay_mbar_ptr.data_ptr(), - ).make_participants() - # TMEM tmem_alloc_barrier = pipeline.NamedBarrier( barrier_id=1, @@ -1212,54 +1261,111 @@ def kernel( sK_g = storage.sK.get_tensor(k_smem_layout_staged.outer, swizzle=k_smem_layout_staged.inner) sK_kv = storage.sQ_K_scaled.get_tensor(kv_k_smem_mma_layout_staged.outer, swizzle=kv_k_smem_mma_layout_staged.inner) sK_ks = storage.sQ_K_scaled.get_tensor(k_smem_mma_layout_staged.outer, swizzle=k_smem_mma_layout_staged.inner) - # NOTE: reuse same smem as sG - sK_neg_g_f32 = storage.sG.get_tensor( - # kv_k_smem_layout_staged.outer, swizzle=kv_k_smem_layout_staged.inner - # NOTE: same swizzle atom (k-major) as k_smem_layout_staged - k_smem_layout_staged.outer, - swizzle=k_smem_layout_staged.inner, - ) - # NOTE: recast as bf16 since operand B is BF16 - # CRITICAL FIX: sK_neg_g's stage stride must match sG's byte stride - # sG (F32) has stage stride = 8192 elements = 32768 bytes - # sK_neg_g (BF16) must have stage stride = 32768 bytes = 16384 BF16 elements - # Original bug: using sK_g.layout which has stage stride = 8192 BF16 elements = 16384 bytes - # This caused sK_neg_g stage 1 to overlap with sG stage 0's second half! - - # Get the base layout from sK_g but double the stage stride - sK_g_outer = sK_g.layout - # Create new layout with corrected stage stride (16384 BF16 elements instead of 8192) - # sK_g layout is: ((64,16),1,(4,2),2):((64,1),0,(16,4096),8192) - # We need: ((64,16),1,(4,2),2):((64,1),0,(16,4096),16384) - sK_neg_g_layout = cute.make_layout( - sK_g_outer.shape, - stride=(*sK_g_outer.stride[:-1], sK_g_outer.stride[-1] * 2), # Double the stage stride - ) - sK_neg_g = cute.make_tensor( - cute.recast_ptr(sK_neg_g_f32.iterator, swizzle_=k_smem_layout_staged.inner, dtype=self.io_dtype), - layout=sK_neg_g_layout, - ) - - # Same fix for sK_neg_g_b - sK_neg_g_b_outer = kv_k_smem_layout_staged.outer - sK_neg_g_b_layout = cute.make_layout( - sK_neg_g_b_outer.shape, - stride=(*sK_neg_g_b_outer.stride[:-1], sK_neg_g_b_outer.stride[-1] * 2), # Double the stage stride - ) + # Gate staging view and the BF16 K*exp(-g) view which reuses the same + # physical buffer after all gate consumers finish. + if cutlass.const_expr(self.scalar_gate): + sG_tma = storage.sG.get_tensor(g_smem_layout_staged.outer, swizzle=g_smem_layout_staged.inner) + scalar_g_stage_elements = Constant.C * ( + 4 + Constant.C // Constant.SC + ) + (Constant.C // Constant.SC) ** 2 + sG = cute.make_tensor( + sG_tma.iterator, + layout=cute.make_layout( + (Constant.C, Constant.D, self.g_stage), + stride=(1, 0, scalar_g_stage_elements), + ), + ) + sG_factor_a = cute.make_tensor( + sG_tma.iterator + Constant.C, + layout=cute.make_layout( + (Constant.C, self.g_stage), + stride=(1, scalar_g_stage_elements), + ), + ) + sG_factor_b = cute.make_tensor( + sG_tma.iterator + 2 * Constant.C, + layout=cute.make_layout( + (Constant.C // Constant.SC, Constant.C, self.g_stage), + stride=(Constant.C, 1, scalar_g_stage_elements), + ), + ) + sG_factor_base = cute.make_tensor( + sG_tma.iterator + + Constant.C * (2 + Constant.C // Constant.SC), + layout=cute.make_layout( + ( + Constant.C // Constant.SC, + Constant.C // Constant.SC, + self.g_stage, + ), + stride=( + Constant.C // Constant.SC, + 1, + scalar_g_stage_elements, + ), + ), + ) + scalar_base_end = ( + Constant.C * (2 + Constant.C // Constant.SC) + + (Constant.C // Constant.SC) ** 2 + ) + sG_main_q = cute.make_tensor( + sG_tma.iterator + scalar_base_end, + layout=cute.make_layout( + (Constant.C, self.g_stage), + stride=(1, scalar_g_stage_elements), + ), + ) + sG_main_k = cute.make_tensor( + sG_tma.iterator + scalar_base_end + Constant.C, + layout=cute.make_layout( + (Constant.C, self.g_stage), + stride=(1, scalar_g_stage_elements), + ), + ) + # The BF16 scratch views below are compile-time dead in the + # supported safe-gate scalar specialization. + sK_neg_g_layout = sK_g.layout + sK_neg_g_b_layout = kv_k_smem_layout_staged.outer + sK_neg_g_ptr = cute.recast_ptr( + sG_tma.iterator, + swizzle_=k_smem_layout_staged.inner, + dtype=self.io_dtype, + ) + else: + sG_tma = storage.sG.get_tensor(g_smem_layout_staged.outer, swizzle=g_smem_layout_staged.inner) + sG = sG_tma + # Compile-time unused by the generic vector-gate branch. + sG_factor_a = sG_tma + sG_factor_b = sG_tma + sG_factor_base = sG_tma + sG_main_q = sG_tma + sG_main_k = sG_tma + # Vector G occupies FP32 stages, so BF16 scratch must preserve the + # corresponding byte stride when it is overlaid on that storage. + sK_g_outer = sK_g.layout + sK_neg_g_layout = cute.make_layout( + sK_g_outer.shape, + stride=(*sK_g_outer.stride[:-1], sK_g_outer.stride[-1] * 2), + ) + sK_neg_g_b_outer = kv_k_smem_layout_staged.outer + sK_neg_g_b_layout = cute.make_layout( + sK_neg_g_b_outer.shape, + stride=(*sK_neg_g_b_outer.stride[:-1], sK_neg_g_b_outer.stride[-1] * 2), + ) + sK_neg_g_ptr = cute.recast_ptr( + sG_tma.iterator, + swizzle_=k_smem_layout_staged.inner, + dtype=self.io_dtype, + ) + + sK_neg_g = cute.make_tensor(sK_neg_g_ptr, layout=sK_neg_g_layout) sK_neg_g_b = cute.make_tensor( - cute.recast_ptr(sK_neg_g_f32.iterator, swizzle_=kv_k_smem_layout_staged.inner, dtype=self.io_dtype), + cute.recast_ptr(sG_tma.iterator, swizzle_=kv_k_smem_layout_staged.inner, dtype=self.io_dtype), layout=sK_neg_g_b_layout, ) - # sK_neg_g = cute.make_tensor( - # cute.recast_ptr( - # sK_neg_g_f32.iterator, - # swizzle_=k_smem_layout_staged.inner, - # dtype=self.io_dtype), - # layout=sK_neg_g_f32.layout) # (((64,2),16),1,4,2):(((1,4096),64),0,1024,8192)> sV = storage.sV.get_tensor(v_smem_layout_staged.outer, swizzle=v_smem_layout_staged.inner) - # G (gate/g_cumsum) - NEW for KDA - sG = storage.sG.get_tensor(g_smem_layout_staged.outer, swizzle=g_smem_layout_staged.inner) # No swizzling for last row of exp(G) sG_last = self.get_smem_tensor_sG_last(storage, g_last_layout) @@ -1315,7 +1421,7 @@ def kernel( self.v_dtype, # utils.LayoutEnum.ROW_MAJOR, utils.LayoutEnum.COL_MAJOR, - (Constant.D, Constant.C), + (self.value_tile_size, Constant.C), self.v_stage, ) v_smem_layout_coalesce = cute.coalesce( @@ -1374,22 +1480,24 @@ def kernel( ) sK_flat_s2r = storage.sK.get_tensor(k_smem_layout_coalesce.outer, swizzle=k_smem_layout_coalesce.inner) - sG_flat_s2r_f32_fake = storage.sG.get_tensor(k_smem_layout_coalesce.outer, swizzle=k_smem_layout_coalesce.inner) - # CRITICAL FIX: When recasting F32 to BF16, we must double the stage stride - # so that the byte offset remains the same. - # F32 stage stride = 8192 elements = 32768 bytes - # BF16 stage stride should = 16384 elements = 32768 bytes k_smem_layout_bf16_outer = k_smem_layout_coalesce.outer - k_smem_layout_bf16_fixed = cute.make_layout( - k_smem_layout_bf16_outer.shape, - stride=(*k_smem_layout_bf16_outer.stride[:-1], k_smem_layout_bf16_outer.stride[-1] * 2), - ) + if cutlass.const_expr(self.scalar_gate): + # Scalar G uses dense BF16 scratch stages. + k_smem_layout_bf16_fixed = k_smem_layout_bf16_outer + else: + # Vector G has 32 KiB FP32 stages; retain their byte stride in the + # BF16 overlay. + k_smem_layout_bf16_fixed = cute.make_layout( + k_smem_layout_bf16_outer.shape, + stride=(*k_smem_layout_bf16_outer.stride[:-1], k_smem_layout_bf16_outer.stride[-1] * 2), + ) sG_flat_bf16 = cute.make_tensor( - cute.recast_ptr(sG_flat_s2r_f32_fake.iterator, swizzle_=k_smem_layout_coalesce.inner, dtype=self.io_dtype), + cute.recast_ptr(sG_tma.iterator, swizzle_=k_smem_layout_coalesce.inner, dtype=self.io_dtype), layout=k_smem_layout_bf16_fixed, ) - (_, hidx, bidx) = cute.arch.block_idx() + (value_tile_idx, hidx, bidx) = cute.arch.block_idx() + value_tile_base = value_tile_idx * self.value_tile_size B, S, H, HV, D = problem_size qk_hidx = hidx // (HV // H) C = self.chunk_size @@ -1514,20 +1622,22 @@ def kernel( # ------------------------------------------------------- # ((SWIZZLE_ATOM_M, REST_M), (SWIZZLE_ATOM_N, REST_N), (1, STAGES)) - g_smem_layout_epi = sm100_utils.make_smem_layout_epi( - self.g_dtype, - utils.LayoutEnum.ROW_MAJOR, - # G SMEM has the shape of - (Constant.C, Constant.D), - self.g_stage, - ) - # (C, (SWIZZLE_ATOM_N, REST_N), STAGES) - g_smem_layout_coalesce = cute.coalesce( - g_smem_layout_epi, - target_profile=(1, 1, 1), - ) - # ROW MAJOR - sG_flat = storage.sG.get_tensor(g_smem_layout_coalesce.outer, swizzle=g_smem_layout_coalesce.inner) + if cutlass.const_expr(self.scalar_gate): + # Logical vector view used by the existing safe-gate arithmetic. + # The D mode has zero stride, so no vector gate is materialized. + sG_flat = sG + else: + g_smem_layout_epi = sm100_utils.make_smem_layout_epi( + self.g_dtype, + utils.LayoutEnum.ROW_MAJOR, + (Constant.C, Constant.D), + self.g_stage, + ) + g_smem_layout_coalesce = cute.coalesce( + g_smem_layout_epi, + target_profile=(1, 1, 1), + ) + sG_flat = storage.sG.get_tensor(g_smem_layout_coalesce.outer, swizzle=g_smem_layout_coalesce.inner) # /////////////////////////////////////////////////////////////////////////////// # LOAD WARP # /////////////////////////////////////////////////////////////////////////////// @@ -1582,17 +1692,17 @@ def kernel( batch_idx=data_bidx, ) - # G (gate) - NEW for KDA - tGsG, tGgG = self.tma_partition_for_mma_operand( - tma_atom_g, - tma_tensor_g_v, - sG, - self.qk_mma_tiler, # Same as Q - qk_tiled_mma, - operand_mode="A", - debug_name="G", - batch_idx=data_bidx, - ) + if cutlass.const_expr(not self.scalar_gate): + tGsG, tGgG = self.tma_partition_for_mma_operand( + tma_atom_g, + tma_tensor_g_v, + sG_tma, + self.qk_mma_tiler, + qk_tiled_mma, + operand_mode="A", + debug_name="G", + batch_idx=data_bidx, + ) if cutlass.const_expr(PRINT_DEBUG): print(f"tKsK={tKsK}") @@ -1605,29 +1715,25 @@ def kernel( idx = chunk_start // C should_debug = PRINT_DEBUG and tidx == warp_idx * 32 and hidx == 0 and bidx == 0 - # Gi (gate/g_cumsum) - NEW for KDA - g_handle = load_g_producer.acquire_and_advance() - cute.copy( - atom=tma_atom_g, - src=tGgG[None, idx, 0], - dst=tGsG[None, g_handle.index], - tma_bar_ptr=g_handle.barrier, - ) + # Vector G remains a TMA operand on the load warp. Scalar G + # is produced independently by load_beta_warp below, allowing + # this warp to issue Q/K/V immediately. + if cutlass.const_expr(not self.scalar_gate): + g_handle = load_g_producer.acquire_and_advance() + cute.copy( + atom=tma_atom_g, + src=tGgG[None, idx, 0], + dst=tGsG[None, g_handle.index], + tma_bar_ptr=g_handle.barrier, + ) - # Qi - # SRC: ((ATOM_V, REST_V), TILES_M, TILES_K) - # DST: ((ATOM_V, REST_V), INPUT_STAGE) q_handle = load_q_producer.acquire_and_advance() cute.copy( atom=tma_atom_q, - src=tQgQ[None, idx, 0], # source - dst=tQsQ[None, q_handle.index], # which stage + src=tQgQ[None, idx, 0], + dst=tQsQ[None, q_handle.index], tma_bar_ptr=q_handle.barrier, ) - - # Ki - # SRC: ((ATOM_V, REST_V), TILES_N, TILES_K) - # DST: ((ATOM_V, REST_V), INPUT_STAGE) k_handle = load_k_producer.acquire_and_advance() cute.copy( atom=tma_atom_k, @@ -1635,16 +1741,10 @@ def kernel( dst=tKsK[None, k_handle.index], tma_bar_ptr=k_handle.barrier, ) - - # Vi - # SRC: ((ATOM_V, REST_V), TILES_M, TILES_K) - # DST: ((ATOM_V, REST_V), INPUT_STAGE) v_handle = load_v_producer.acquire_and_advance() - if cutlass.const_expr(PRINT_DEBUG) and should_debug: - cute.printf("TMA v producer idx={}, v_handle={}", idx, v_handle.index) cute.copy( atom=tma_atom_v, - src=tVgV[None, 0, idx], + src=tVgV[None, value_tile_idx, idx], dst=tVsV[None, v_handle.index], tma_bar_ptr=v_handle.barrier, ) @@ -1678,7 +1778,7 @@ def kernel( tCtAcc=tCtAccSQ, tCrA=tCrState, tCrB=tCrQ_sq, - a_stage_idx=0, + a_stage_idx=kv16_handle.index, b_stage_idx=q_scaled_handle.index, acc_stage_idx=0, ) @@ -1751,7 +1851,6 @@ def kernel( # wait for sQ_K_scaled and decay(S) ready k_scaled2_handle = load_k_scaled2_consumer.wait_and_advance() - kv_decay_handle = kv_decay_consumer.wait_and_advance() kv_handle = kv_producer.acquire_and_advance() # launch S=K^T@NewV MMA @@ -1769,7 +1868,6 @@ def kernel( ) k_scaled2_handle.release() - kv_decay_handle.release() kv_handle.commit() # NOTE: Add a signal to notify the end of v consumption. @@ -1847,7 +1945,7 @@ def kernel( tCtAcc=tCtAccSQ, tCrA=tCrState, tCrB=tCrQ_sq, - a_stage_idx=0, + a_stage_idx=kv16_handle.index, b_stage_idx=q_handle.index, acc_stage_idx=0, ) @@ -2135,7 +2233,7 @@ def kernel( tCtAccKV_slice = tCtAccKV[((None, None), 0, 0, None)] ( tiled_copy_t2r_kv, - _, # thr_t2r + thr_t2r_kv, tTR_tKV, tTR_rKV, ) = self.tmem_load_partition_kv( @@ -2143,6 +2241,9 @@ def kernel( tState=tCtAccKV_slice, local_tidx=local_tidx, ) + state_tile = cute.dice(self.kv_mma_tiler, (1, 1, None)) + cM_state = cute.make_identity_tensor(state_tile) + tTR_cState = thr_t2r_kv.partition_D(cM_state) ############################################################ ( @@ -2155,15 +2256,35 @@ def kernel( ) tmem_store_rKV = cute.make_tensor(tTR_rKV.iterator, layout=tmem_store_rAccKV_f32.layout) - ( - tmem_store_kv, - tmem_store_tAccKV, - tmem_store_rAccKV, - ) = self.tmem_store_and_partition_acc( - local_tidx, - tCtAcc=tCtStateAsF32, - ) - tmem_store_rAccKVAsBF16 = cute.recast_tensor(tmem_store_rAccKV, dtype=self.io_dtype) + if cutlass.const_expr(self.split_value_tiles): + # For M=64, publish the BF16 state through the native operand-A + # TMEM layout. This is the same mapping used by the proven + # chunk_delta_h SM100 kernel; the original M=128 packed-FP32 + # alias has a different TV ownership and scrambles state reads + # in the following chunk. + state_store_atom = cute.make_copy_atom( + tcgen05.St16x128bOp(tcgen05.Repetition(16), tcgen05.Unpack.NONE), + self.io_dtype, + ) + tmem_store_kv = tcgen05.make_tmem_copy(state_store_atom, tCrState) + state_store_thr = tmem_store_kv.get_slice(local_tidx) + state_store_shape = cute.slice_( + state_store_thr.partition_S(tCrState).shape, + (None, None, None, None, 0), + ) + tmem_store_tAccKV = state_store_thr.partition_D(tCrState) + tmem_store_rAccKV = cute.make_rmem_tensor(state_store_shape, self.io_dtype) + tmem_store_rAccKVAsBF16 = tmem_store_rAccKV + else: + ( + tmem_store_kv, + tmem_store_tAccKV, + tmem_store_rAccKV, + ) = self.tmem_store_and_partition_acc( + local_tidx, + tCtAcc=tCtStateAsF32, + ) + tmem_store_rAccKVAsBF16 = cute.recast_tensor(tmem_store_rAccKV, dtype=self.io_dtype) ############################################################ if cutlass.const_expr(PRINT_DEBUG): @@ -2256,7 +2377,7 @@ def kernel( # ------------------------------------------------------- # V s2r partitions - NEW for KDA elementwise processing # shape_v = (Constant.C, Constant.D) - shape_v = (Constant.D, Constant.C) + shape_v = (self.value_tile_size, Constant.C) ( tiled_s2r_v, thr_s2r_v, @@ -2283,8 +2404,18 @@ def kernel( thr_mma_epi_half = tiled_mma_epi_half.get_slice(local_tidx) copy_op_qk_s2r = cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4) copy_op_qk_r2s = cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=False, num_matrices=4) - # FIXME: only 2 FP32 elements (64 bits) compatible with ldmatrix, how to change to 128? - copy_atom_g = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.g_dtype, num_bits_per_copy=64) + # A scalar-gate broadcast has a zero-stride D mode; use scalar + # copies so every logical element observes that stride. The + # generic vector path retains its 64-bit copy atom. + if cutlass.const_expr(self.scalar_gate): + gate_copy_bits = 32 + else: + gate_copy_bits = 64 + copy_atom_g = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.g_dtype, + num_bits_per_copy=gate_copy_bits, + ) # Half-size tiled copies for partitioned S2R tiled_load_g_half = cute.make_tiled_copy_A(copy_atom_g, tiled_mma_epi_half) thr_load_g_half = tiled_load_g_half.get_slice(local_tidx) @@ -2332,19 +2463,29 @@ def kernel( sK_flat_h0 = cute.make_tensor(sK_flat_s2r.iterator, layout=k_half_outer) sK_flat_h1 = cute.make_tensor(sK_flat_s2r.iterator + HALF_SMEM_ELEMS, layout=k_half_outer) - # G half SMEM views (FP32, from sG_flat iterator) - g_sml_epi_half = sm100_utils.make_smem_layout_epi( - self.g_dtype, - utils.LayoutEnum.ROW_MAJOR, - (Constant.C, Constant.HALF_D), - self.g_stage, - ) - g_sml_half = cute.coalesce(g_sml_epi_half, target_profile=(1, 1, 1)) - g_half_outer = cute.make_layout( - g_sml_half.outer.shape, stride=(*g_sml_half.outer.stride[:-1], g_sml_half.outer.stride[-1] * 2) - ) - sG_flat_h0 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) - sG_flat_h1 = cute.make_tensor(sG_flat.iterator + HALF_SMEM_ELEMS, layout=g_half_outer) + # G half SMEM views (FP32, from sG_flat iterator). + if cutlass.const_expr(self.scalar_gate): + g_half_outer = cute.make_layout( + (Constant.C, Constant.HALF_D, self.g_stage), + stride=(1, 0, Constant.C), + ) + sG_flat_h0 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) + # Both D halves broadcast the same per-token scalar. + sG_flat_h1 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) + else: + g_sml_epi_half = sm100_utils.make_smem_layout_epi( + self.g_dtype, + utils.LayoutEnum.ROW_MAJOR, + (Constant.C, Constant.HALF_D), + self.g_stage, + ) + g_sml_half = cute.coalesce(g_sml_epi_half, target_profile=(1, 1, 1)) + g_half_outer = cute.make_layout( + g_sml_half.outer.shape, + stride=(*g_sml_half.outer.stride[:-1], g_sml_half.outer.stride[-1] * 2), + ) + sG_flat_h0 = cute.make_tensor(sG_flat.iterator, layout=g_half_outer) + sG_flat_h1 = cute.make_tensor(sG_flat.iterator + HALF_SMEM_ELEMS, layout=g_half_outer) # Q_K_scaled half SMEM views (from sQ_K_scaled_flat iterator) qks_sml_epi_half = sm100_utils.make_smem_layout_epi( @@ -2405,12 +2546,26 @@ def index_transform_half(index_q, index_k): # State shape: (D, D) per (H, B), stored as FP32 if cutlass.const_expr(self.has_initial_state): # Load initial state from GMEM to RMEM respecting TMEM partition. - # TMEM stores S^T (transposed), so flat[i] = state[local_tidx, i] - # Each thread owns key position local_tidx, D elements cover value positions. init_state_chunk = initial_state[None, None, (hidx, bidx)] - init_flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) - for init_i in cutlass.range(0, Constant.D, unroll=0): - init_flat[init_i] = init_state_chunk[local_tidx, init_i] + if cutlass.const_expr(self.split_value_tiles): + # M=64 uses 16 TMEM datapaths/warp, so thread id is no + # longer the K coordinate. Follow the actual T2R TV map. + for init_i in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + value_coord, key_coord = tTR_cState[init_i] + tTR_rKV[init_i] = init_state_chunk[ + key_coord, + value_tile_base + value_coord, + ] + else: + init_flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) + for init_i in cutlass.range(0, self.value_tile_size, unroll=0): + init_flat[init_i] = init_state_chunk[ + local_tidx, + value_tile_base + init_i, + ] # Store FP32 state to TMEM for accumulation (tCtAccKV) init_tmem_store_tKVi = tmem_store_tAccKV_f32[None, None, None, None, 0] @@ -2430,9 +2585,8 @@ def index_transform_half(index_q, index_k): # safe_gate version for chunk_start in cutlass.range(0, seq_len, C, unroll=0): idx = chunk_start // C - if cutlass.const_expr(self.is_varlen): - valid_len_chunk = seq_len - chunk_start - else: + valid_len_chunk = seq_len - chunk_start + if valid_len_chunk > C: valid_len_chunk = C # ============================================================ @@ -2462,6 +2616,14 @@ def index_transform_half(index_q, index_k): for half_idx in cutlass.range_constexpr(2): tQrG_persists.append(cute.make_fragment_like(tQrQ_half_0, dtype=self.g_dtype)) + if cutlass.const_expr(self.scalar_gate): + if local_tidx == 0: + sG_last[0, g_stage_idx] = sG_tma[ + valid_len_chunk - 1, + 0, + g_stage_idx, + ] + # Merged g_last + Q gating path: single G half-load per half if idx != 0 or cutlass.const_expr(self.has_initial_state): q_stage_idx = q_handle.index @@ -2473,38 +2635,50 @@ def index_transform_half(index_q, index_k): # S2R G half into persistent fragment tQrG_half_cv = thr_load_g_half.retile(tQrG_persists[half_idx]) - cute.copy(tiled_load_g_half, tQsG_h[half_idx][None, None, None, g_stage_idx], tQrG_half_cv) - - # Write g_last half (before exp transforms g values) - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if cutlass.const_expr(self.is_varlen): + if cutlass.const_expr(self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, _ = index_transform_half(*tQcMq_half[i]) + tQrG_persists[half_idx][i] = sG_main_q[index_q, g_stage_idx] + else: + cute.copy( + tiled_load_g_half, + tQsG_h[half_idx][None, None, None, g_stage_idx], + tQrG_half_cv, + ) + + # Vector G needs one g_last value per feature; + # scalar G wrote its single value above. + if cutlass.const_expr(not self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) if valid_len_chunk < C: if index_q == valid_len_chunk - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] else: if index_q == Constant.C - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] - else: - if index_q == Constant.C - 1: - sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] - # exp(g) half in-place — persists for K gating reuse - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - tQrG_persists[half_idx][i] = cute.exp2(tQrG_persists[half_idx][i], fastmath=self.use_fast_math) + # exp(g) half in-place — persists for K gating reuse. + # Scalar G was already exponentiated once/token by + # the producer warp, rather than once/feature here. + if cutlass.const_expr(not self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + tQrG_persists[half_idx][i] = cute.exp2( + tQrG_persists[half_idx][i], + fastmath=self.use_fast_math, + ) # S2R Q half tQrQ_half = cute.make_fragment_like(tQrQ_half_0, self.q_dtype) tQrQ_half_cv = thr_load_qk_half.retile(tQrQ_half) cute.copy(tiled_load_qk_half, tQsQ_h[half_idx][None, None, None, q_stage_idx], tQrQ_half_cv) - # Zero Q for invalid positions (varlen only) - if cutlass.const_expr(self.is_varlen): + # Zero Q for any partial tail (varlen or fixed-B). + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) if valid_len_chunk < C: - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if index_q >= valid_len_chunk: - tQrQ_half[i] = self.q_dtype(0.0) + if index_q >= valid_len_chunk: + tQrQ_half[i] = self.q_dtype(0.0) # Q gating: Q' = Q * exp(g) * scale for i in cutlass.range_constexpr(cute.size(tQcMq_half)): @@ -2527,19 +2701,21 @@ def index_transform_half(index_q, index_k): for half_idx in cutlass.range_constexpr(2): k_offset = half_idx * Constant.HALF_D tQrG_half_cv = thr_load_g_half.retile(tQrG_persists[half_idx]) - cute.copy(tiled_load_g_half, tQsG_h[half_idx][None, None, None, g_stage_idx], tQrG_half_cv) - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if cutlass.const_expr(self.is_varlen): + if cutlass.const_expr(not self.scalar_gate): + cute.copy( + tiled_load_g_half, + tQsG_h[half_idx][None, None, None, g_stage_idx], + tQrG_half_cv, + ) + if cutlass.const_expr(not self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) if valid_len_chunk < C: if index_q == valid_len_chunk - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] else: if index_q == Constant.C - 1: sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] - else: - if index_q == Constant.C - 1: - sG_last[index_k + k_offset, g_stage_idx] = tQrG_persists[half_idx][i] # ==================================================== # Partitioned S2R: K gating (2 half-passes) @@ -2564,13 +2740,12 @@ def index_transform_half(index_q, index_k): tQrK_half_cv = thr_load_qk_half.retile(tQrK_half) cute.copy(tiled_load_qk_half, tQsK_h[half_idx][None, None, None, k_stage_idx], tQrK_half_cv) - # Zero K for invalid positions (varlen only) - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C: - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if index_q >= valid_len_chunk: - tQrK_half[i] = self.q_dtype(0.0) + # Zero K for any partial tail (varlen or fixed-B). + if valid_len_chunk < C: + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) + if index_q >= valid_len_chunk: + tQrK_half[i] = self.q_dtype(0.0) # K gating: K' = K * exp(g) — reuse persisted exp2(g) for i in cutlass.range_constexpr(cute.size(tQcMq_half)): @@ -2693,7 +2868,6 @@ def index_transform_half(index_q, index_k): # decay S, S=S*g_last # FIXME: currently do not support initial state, # so only decay S after first block K^T@NewV - kv_decay_handle = kv_decay_producer.acquire_and_advance() if idx != 0 or cutlass.const_expr(self.has_initial_state): # NOTE: TMEM S is always ready here # T2R S @@ -2702,7 +2876,10 @@ def index_transform_half(index_q, index_k): cute.arch.fence_view_async_tmem_load() # decay S - flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) + flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) self.scale_state(flat, sG_last[None, g_stage_idx]) # R2T S @@ -2711,8 +2888,6 @@ def index_transform_half(index_q, index_k): cute.copy(tmem_store_kv_f32, tmem_store_rKV, tmem_store_tKVi) cute.arch.fence_view_async_tmem_store() - kv_decay_handle.commit() - # ==================================================== # Partitioned S2R: K^T gating — exp(g_last-g)*K # ==================================================== @@ -2725,20 +2900,28 @@ def index_transform_half(index_q, index_k): # S2R G half tQrG_half = cute.make_fragment_like(tQrQ_half_0, dtype=self.g_dtype) tQrG_half_cv = thr_load_g_half.retile(tQrG_half) - cute.copy(tiled_load_g_half, tQsG_h[half_idx][None, None, None, g_stage_idx], tQrG_half_cv) + if cutlass.const_expr(self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, _ = index_transform_half(*tQcMq_half[i]) + tQrG_half[i] = sG_main_k[index_q, g_stage_idx] + else: + cute.copy( + tiled_load_g_half, + tQsG_h[half_idx][None, None, None, g_stage_idx], + tQrG_half_cv, + ) # S2R K half tQrK_half = cute.make_fragment_like(tQrQ_half_0, dtype=self.k_dtype) tQrK_half_cv = thr_load_qk_half.retile(tQrK_half) cute.copy(tiled_load_qk_half, tQsK_h[half_idx][None, None, None, k_stage_idx], tQrK_half_cv) - # Zero K half for invalid positions (varlen only) - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C: - for i in cutlass.range_constexpr(cute.size(tQcMq_half)): - index_q, index_k = index_transform_half(*tQcMq_half[i]) - if index_q >= valid_len_chunk: - tQrK_half[i] = self.k_dtype(0.0) + # Zero K half for any partial tail. + if valid_len_chunk < C: + for i in cutlass.range_constexpr(cute.size(tQcMq_half)): + index_q, index_k = index_transform_half(*tQcMq_half[i]) + if index_q >= valid_len_chunk: + tQrK_half[i] = self.k_dtype(0.0) # K^T gating: exp(g_last - g) * K for i in cutlass.range_constexpr(cute.size(tQcMq_half)): @@ -2746,7 +2929,14 @@ def index_transform_half(index_q, index_k): g_last_val = sG_last[index_k + k_offset, g_stage_idx] k_i = tQrK_half[i].to(cutlass.Float32) g_i = tQrG_half[i] - tQrK_half[i] = (cute.exp2(g_last_val - g_i, fastmath=self.use_fast_math) * k_i).to(self.k_dtype) + if cutlass.const_expr(self.scalar_gate): + gate_k_i = g_i + else: + gate_k_i = cute.exp2( + g_last_val - g_i, + fastmath=self.use_fast_math, + ) + tQrK_half[i] = (gate_k_i * k_i).to(self.k_dtype) # R2S K^T half to sQ_K_scaled tQrK_half_cv_src = thr_store_qk_half.retile(tQrK_half) @@ -2815,6 +3005,7 @@ def index_transform_half(index_q, index_k): cute.copy(tiled_copy_t2r_kv, tTR_tKVi, tTR_rKV) cute.arch.fence_view_async_tmem_load() + if idx != final_blk: # Store as a separated BF16 state for QS and KS MMA before decay # tmem_store_rAccKVAsBF16 point to the same rmem as tmem_store_rKV @@ -2836,11 +3027,28 @@ def index_transform_half(index_q, index_k): # idx == final_blk: output final state immediately to minimize tTR_rKV lifetime if cutlass.const_expr(self.output_final_state): # Write FP32 state from RMEM to GMEM - # TMEM stores S^T (transposed), so flat[i] = state[local_tidx, i] state_out = final_state[None, None, (hidx, bidx)] - out_flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) - for out_i in cutlass.range(0, Constant.D, unroll=0): - state_out[local_tidx, out_i] = out_flat[out_i] + if cutlass.const_expr(self.split_value_tiles): + # Core ABI: h0 is contiguous [V,K], while ht is + # returned contiguous [K,V]. With fstate's + # CuTe stride (1,D), storing (V,K) below creates + # the required row-major [K,V] result. + for out_i in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + value_coord, key_coord = tTR_cState[out_i] + state_out[ + value_tile_base + value_coord, + key_coord, + ] = tTR_rKV[out_i] + else: + out_flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) + for out_i in cutlass.range(0, self.value_tile_size, unroll=0): + state_out[ + local_tidx, + value_tile_base + out_i, + ] = out_flat[out_i] # release KV kv_handle.release() @@ -2855,9 +3063,8 @@ def index_transform_half(index_q, index_k): else: for chunk_start in cutlass.range(0, seq_len, C, unroll=0): idx = chunk_start // C - if cutlass.const_expr(self.is_varlen): - valid_len_chunk = seq_len - chunk_start - else: + valid_len_chunk = seq_len - chunk_start + if valid_len_chunk > C: valid_len_chunk = C # ============================================================ @@ -2924,19 +3131,10 @@ def index_transform_half(index_q, index_k): # Element-wise processing avoids bulk .load()/.to() creating # ~300+ register SSA vectors from G, Q, K simultaneously for _zr in cutlass.range(0, Constant.C, unroll_full=True): - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C and _zr >= valid_len_chunk: - tRS_rQ[0, _zr, 0] = self.io_dtype(0.0) - tRS_rK[0, _zr, 0] = self.io_dtype(0.0) - tRS_rG_bf16[0, _zr, 0] = self.io_dtype(0.0) - else: - g_i = tRS_rG[0, _zr, 0] - exp_g_i = cute.exp2(g_i, fastmath=self.use_fast_math) - q_i = tRS_rQ[0, _zr, 0].to(cutlass.Float32) - tRS_rQ[0, _zr, 0] = (q_i * exp_g_i * self.scale).to(self.io_dtype) - k_i = tRS_rK[0, _zr, 0].to(cutlass.Float32) - tRS_rK[0, _zr, 0] = (k_i * exp_g_i).to(self.io_dtype) - tRS_rG_bf16[0, _zr, 0] = (k_i * cute.exp2(-g_i, fastmath=self.use_fast_math)).to(self.io_dtype) + if valid_len_chunk < C and _zr >= valid_len_chunk: + tRS_rQ[0, _zr, 0] = self.io_dtype(0.0) + tRS_rK[0, _zr, 0] = self.io_dtype(0.0) + tRS_rG_bf16[0, _zr, 0] = self.io_dtype(0.0) else: g_i = tRS_rG[0, _zr, 0] exp_g_i = cute.exp2(g_i, fastmath=self.use_fast_math) @@ -3011,13 +3209,10 @@ def index_transform_half(index_q, index_k): # NOTE: Save exp(g) of last VALID row to rG_last for state update in next chunk # For full chunks, directly use C-1; only loop for partial chunks (varlen only) rG_last = cutlass.Float32(1.0) - if cutlass.const_expr(self.is_varlen): - if valid_len_chunk < C: - for _zr in cutlass.range(0, Constant.C, unroll_full=True): - if _zr == valid_len_chunk - 1: - rG_last = cute.exp2(tRS_rG[0, _zr, 0], fastmath=self.use_fast_math) - else: - rG_last = cute.exp2(tRS_rG[0, Constant.C - 1, 0], fastmath=self.use_fast_math) + if valid_len_chunk < C: + for _zr in cutlass.range(0, Constant.C, unroll_full=True): + if _zr == valid_len_chunk - 1: + rG_last = cute.exp2(tRS_rG[0, _zr, 0], fastmath=self.use_fast_math) else: rG_last = cute.exp2(tRS_rG[0, Constant.C - 1, 0], fastmath=self.use_fast_math) # NOTE: each thread save one element @@ -3318,7 +3513,10 @@ def index_transform_half(index_q, index_k): cute.print_tensor(tTR_rKV) self.cuda_wg_sync_barrier.arrive_and_wait() - flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) + flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) # FIXME self.cuda_wg_sync_barrier.arrive_and_wait() @@ -3360,9 +3558,25 @@ def index_transform_half(index_q, index_k): # idx == final: output final state immediately to minimize tTR_rKV lifetime if cutlass.const_expr(self.output_final_state): state_out = final_state[None, None, (hidx, bidx)] - out_flat = cute.make_tensor(tTR_rKV.iterator, layout=cute.make_layout(Constant.D)) - for out_i in cutlass.range(0, Constant.D, unroll=0): - state_out[local_tidx, out_i] = out_flat[out_i] + if cutlass.const_expr(self.split_value_tiles): + # See the safe-gate branch above for the state + # ABI and CuTe-to-row-major axis convention. + for out_i in cutlass.range(cute.size(tTR_rKV), unroll_full=True): + value_coord, key_coord = tTR_cState[out_i] + state_out[ + value_tile_base + value_coord, + key_coord, + ] = tTR_rKV[out_i] + else: + out_flat = cute.make_tensor( + tTR_rKV.iterator, + layout=cute.make_layout(self.value_tile_size), + ) + for out_i in cutlass.range(0, self.value_tile_size, unroll=0): + state_out[ + local_tidx, + value_tile_base + out_i, + ] = out_flat[out_i] # NOTE: only release v after PV and State=KV has been consumed end_v_handle = end_v_consumer.wait_and_advance() @@ -3391,8 +3605,15 @@ def index_transform_half(index_q, index_k): copy_op_A_s2r = cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4) copy_op_B_s2r = cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4) copy_op_r2s = cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2) - # FIXME: only 2 FP32 elements (64 bits) compatible with ldmatrix, how to change to 128? - copy_g_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.g_dtype, num_bits_per_copy=64) + if cutlass.const_expr(self.scalar_gate): + gate_copy_bits = 32 + else: + gate_copy_bits = 64 + copy_g_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.g_dtype, + num_bits_per_copy=gate_copy_bits, + ) G_Q_tiled_copy = cute.make_tiled_copy_A(copy_g_atom, tiled_mma_subchunk) G_Kt_tiled_copy = cute.make_tiled_copy_B(copy_g_atom, tiled_mma_subchunk) Q_tiled_copy = cute.make_tiled_copy_A(cute.make_copy_atom(copy_op_A_s2r, self.q_dtype), tiled_mma_subchunk) @@ -3411,6 +3632,9 @@ def index_transform_half(index_q, index_k): # index tensor cMqk_subchunk = cute.make_identity_tensor(self.qk_kk_subchunk_mma_tiler[:2]) tQKcMqk_subchunk = thr_mma_subchunk.partition_C(cMqk_subchunk) + cA_subchunk = cute.make_identity_tensor((Constant.SC, Constant.BK_SC)) + tAcA_subchunk = thr_mma_subchunk.partition_A(cA_subchunk) + tBcB_subchunk = thr_mma_subchunk.partition_B(cA_subchunk) def index_transform(index_q, index_k): return ( @@ -3459,13 +3683,39 @@ def index_transform(index_q, index_k): sQqk_curr = sQ_flat[None, None, load_q_consumer._PipelineConsumer__state.index] sKqk_curr = sK_flat[None, None, load_k_consumer._PipelineConsumer__state.index] sGqkq_curr = sG_flat[None, None, load_g_consumer._PipelineConsumer__state.index] + if cutlass.const_expr(self.scalar_gate): + sG_scalar_curr = sG_tma[None, 0, load_g_consumer._PipelineConsumer__state.index] + sG_factor_a_curr = sG_factor_a[None, load_g_consumer._PipelineConsumer__state.index] + sG_factor_b_curr = sG_factor_b[None, None, load_g_consumer._PipelineConsumer__state.index] + else: + # The helper signatures are shared with the scalar + # specialization, but this argument is compile-time + # unused by the vector-gate branch. Keep a congruent + # tensor here instead of indexing vector sG as scalar. + sG_scalar_curr = sGqkq_curr + sG_factor_a_curr = sGqkq_curr + sG_factor_b_curr = sGqkq_curr sBeta_curr = sBeta[None, load_beta_consumer._PipelineConsumer__state.index] # (_16,(_32,_2),_4,(_1,_2)):(_32,(_1,_2048),_512,(_0,_4096)) sQqk_slice = cute.flat_divide(sQqk_curr, tiler_subchunk_qk) sKqk_slice = cute.flat_divide(sKqk_curr, tiler_subchunk_qk) # (_16,(_64,_1),_4,(_1,_2)):(_64,(_1,_0),_1024,(_0,_4096)) - sGqkq_slice = cute.flat_divide(sGqkq_curr, tiler_subchunk_g) + if cutlass.const_expr(self.scalar_gate): + # `flat_divide` requires an injective source layout and + # therefore cannot divide the zero-stride broadcast + # view. Build the same logical subchunk coordinates + # directly: token = row + 16*subchunk, while every D + # coordinate aliases that token's scalar gate. + sGqkq_slice = cute.make_tensor( + sGqkq_curr.iterator, + layout=cute.make_layout( + (16, (64, 1), 4, (1, 2)), + stride=(1, (0, 0), 16, (0, 0)), + ), + ) + else: + sGqkq_slice = cute.flat_divide(sGqkq_curr, tiler_subchunk_g) sBeta_slice = cute.flat_divide(sBeta_curr, tiler_subchunk_beta) # Acc results @@ -3523,8 +3773,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3535,7 +3788,10 @@ def index_transform(index_q, index_k): tGsGfirst_0_j_kt = G_Kt_thr_copy.partition_S(sG_first_0_j) tGrGfirst_0_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_0_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_0_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_0_j_kt, tGrGfirst_0_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_0_j_kt.fill(sG_scalar_curr[0]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_0_j_kt, tGrGfirst_0_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3545,7 +3801,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 0, sKqk_slice, tGrGfirst_0_j_kt, ) @@ -3615,8 +3875,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3627,7 +3890,10 @@ def index_transform(index_q, index_k): tGsGfirst_3_j_kt = G_Kt_thr_copy.partition_S(sG_first_3_j) tGrGfirst_3_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_3_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_3_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_3_j_kt, tGrGfirst_3_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_3_j_kt.fill(sG_scalar_curr[3 * Constant.SC]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_3_j_kt, tGrGfirst_3_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3637,7 +3903,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3654,7 +3924,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3671,7 +3945,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3688,7 +3966,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 3, sKqk_slice, tGrGfirst_3_j_kt, ) @@ -3774,8 +4056,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3786,7 +4071,10 @@ def index_transform(index_q, index_k): tGsGfirst_1_j_kt = G_Kt_thr_copy.partition_S(sG_first_1_j) tGrGfirst_1_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_1_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_1_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_1_j_kt, tGrGfirst_1_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_1_j_kt.fill(sG_scalar_curr[Constant.SC]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_1_j_kt, tGrGfirst_1_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3796,7 +4084,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 1, sKqk_slice, tGrGfirst_1_j_kt, ) @@ -3813,7 +4105,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 1, sKqk_slice, tGrGfirst_1_j_kt, ) @@ -3885,8 +4181,11 @@ def index_transform(index_q, index_k): Q_tiled_copy, Q_thr_copy, tv_layout_mma_A, + tAcA_subchunk, layout_g_first, sGqkq_slice, + sG_scalar_curr, + sG_factor_a_curr, sQqk_slice, sKqk_slice, ) @@ -3897,7 +4196,10 @@ def index_transform(index_q, index_k): tGsGfirst_2_j_kt = G_Kt_thr_copy.partition_S(sG_first_2_j) tGrGfirst_2_j_kt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) tGrGfirst_2_j_kt_cv = G_Kt_thr_copy.retile(tGrGfirst_2_j_kt) - cute.copy(G_Kt_tiled_copy, tGsGfirst_2_j_kt, tGrGfirst_2_j_kt_cv) + if cutlass.const_expr(self.scalar_gate): + tGrGfirst_2_j_kt.fill(sG_scalar_curr[2 * Constant.SC]) + else: + cute.copy(G_Kt_tiled_copy, tGsGfirst_2_j_kt, tGrGfirst_2_j_kt_cv) tQKrKt_0_j = self.s2r_compute_subchunk_operand_B( 0, @@ -3907,7 +4209,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 2, sKqk_slice, tGrGfirst_2_j_kt, ) @@ -3924,7 +4230,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 2, sKqk_slice, tGrGfirst_2_j_kt, ) @@ -3941,7 +4251,11 @@ def index_transform(index_q, index_k): Kt_tiled_copy, Kt_thr_copy, tv_layout_mma_B, + tBcB_subchunk, sGqkq_slice, + sG_scalar_curr, + sG_factor_b_curr, + 2, sKqk_slice, tGrGfirst_2_j_kt, ) @@ -3997,9 +4311,8 @@ def index_transform(index_q, index_k): cute.copy(tiled_load_qk, thr_load_qk.partition_S(sQK_curr), tQKrQK) cute.copy(tiled_load_kk, thr_load_kk.partition_S(sKK_inv_curr), tKKrKK) # triangular mask and boundary mask - if cutlass.const_expr(self.is_varlen): - valid_len_chunk = seq_len - chunk_start - else: + valid_len_chunk = seq_len - chunk_start + if valid_len_chunk > C: valid_len_chunk = C self.apply_qk_kk_mask(tQKcMqk, tQKrQK, tKKrKK, valid_len_chunk) # R2S QK/KK @@ -4027,7 +4340,7 @@ def index_transform(index_q, index_k): # S2R, scale with beta, convert to BF16, store back to smem `sM` # TODO: make a repro for the cutedsl team - self.scale_M_inverse_with_beta(local_tidx, sBeta, curr_sM_f16, curr_sM) + self.scale_M_inverse_with_beta(local_tidx, sBeta_curr, curr_sM_f16, curr_sM) cute.arch.fence_proxy( cute.arch.ProxyKind.async_shared, @@ -4138,13 +4451,21 @@ def index_transform(index_q, index_k): cute.copy( tma_atom_o, bSG_sO[None, smem_o_handle.index], - bSG_gO[(None, 0, 0, 0, idx)], + bSG_gO[(None, 0, 0, value_tile_idx, idx)], tma_desc_ptr=self._tensormap_mgr.get_tensormap_ptr(ws_desc_ptr, cute.AddressSpace.generic), ) else: - cute.copy(tma_atom_o, bSG_sO[None, smem_o_handle.index], bSG_gO[(None, 0, 0, 0, idx)]) + cute.copy( + tma_atom_o, + bSG_sO[None, smem_o_handle.index], + bSG_gO[(None, 0, 0, value_tile_idx, idx)], + ) else: - cute.copy(tma_atom_o, bSG_sO[None, smem_o_handle.index], bSG_gO[(None, 0, 0, 0, idx)]) + cute.copy( + tma_atom_o, + bSG_sO[None, smem_o_handle.index], + bSG_gO[(None, 0, 0, value_tile_idx, idx)], + ) # Ensure smem_o has been released. cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0, read=True) @@ -4155,6 +4476,30 @@ def index_transform(index_q, index_k): local_tidx = tidx % (self.threads_per_warp * len([self.load_beta_warp_id])) for chunk_start in cutlass.range(0, seq_len, C, unroll=0): idx = chunk_start // C + if cutlass.const_expr(self.scalar_gate): + g_handle = load_g_producer.acquire_and_advance() + self.produce_scalar_gate_chunk( + g, + sG_tma, + sG_factor_a, + sG_factor_b, + sG_factor_base, + sG_main_q, + sG_main_k, + chunk_start, + seq_len, + tok_offset, + hidx, + data_bidx, + local_tidx, + g_handle.index, + ) + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + g_handle.commit() + # Load beta into smem beta_handle = load_beta_producer.acquire_and_advance() # Fence due to normal load @@ -4180,12 +4525,12 @@ def index_transform(index_q, index_k): valid_len_beta = seq_len - chunk_start for data_idx in cutlass.range(local_tidx, Constant.C, self.threads_per_warp): if data_idx < valid_len_beta: - sBeta[data_idx, 0] = beta_chunk[data_idx, 0] + sBeta[data_idx, beta_handle.index] = beta_chunk[data_idx, 0] else: - sBeta[data_idx, 0] = cutlass.Float32(0.0) + sBeta[data_idx, beta_handle.index] = cutlass.Float32(0.0) else: for data_idx in cutlass.range(local_tidx, Constant.C, self.threads_per_warp): - sBeta[data_idx, 0] = beta_chunk[data_idx, 0] + sBeta[data_idx, beta_handle.index] = beta_chunk[data_idx, 0] # Fence cute.arch.fence_proxy( @@ -4210,12 +4555,19 @@ def scale_state(self, flat: cute.Tensor, sG_last: cute.Tensor) -> cute.Tensor: kv_f32 = flat if cutlass.const_expr(PRINT_DEBUG): print(f"kv_f32: {kv_f32}") - for i in cutlass.range(0, Constant.D, unroll_full=True): - if not cutlass.const_expr(self.safe_gate): - kv_f32[i] = kv_f32[i] * sG_last[i] - else: - # NOTE: when safe_gate=True, sG_last stores the original G values - kv_f32[i] = kv_f32[i] * cute.exp2(sG_last[i], fastmath=self.use_fast_math) + if cutlass.const_expr(self.safe_gate and self.scalar_gate): + # Qwen GDN broadcasts one scalar decay over the full K dimension. + # Compute exp2 once per thread instead of repeating it D times. + decay = cute.exp2(sG_last[0], fastmath=self.use_fast_math) + for i in cutlass.range(0, self.value_tile_size, unroll_full=True): + kv_f32[i] = kv_f32[i] * decay + else: + for i in cutlass.range(0, self.value_tile_size, unroll_full=True): + if not cutlass.const_expr(self.safe_gate): + kv_f32[i] = kv_f32[i] * sG_last[i] + else: + # NOTE: when safe_gate=True, sG_last stores the original G values + kv_f32[i] = kv_f32[i] * cute.exp2(sG_last[i], fastmath=self.use_fast_math) return kv_f32 def tmem_load_kv16(self, local_tidx, tState): @@ -4340,12 +4692,23 @@ def tmem_load_partition_kv(self, mma_tiler, tState, local_tidx): # use_2cta_instrs=False, # ) - # In KDA, we need to make tv-layout row-wise to perform diagonal op. - copy_atom_t2r = cute.make_copy_atom( - # 32b x 32, TODO: ADJUST RMEM PEAK - tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), - self.acc_dtype, - ) + # In KDA, we need a row-wise TV layout for the state operations. The + # original M=128 tile has 32 datapaths per warp, while the split M=64 + # tile has 16; select the matching 16dp load atom for the latter. + if cutlass.const_expr(self.split_value_tiles): + copy_atom_t2r = sm100_utils.get_tmem_load_op( + mma_tiler, + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.acc_dtype, + mma_tiler[:2], + use_2cta_instrs=False, + ) + else: + copy_atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), + self.acc_dtype, + ) fake_sState = cute.make_tensor( cute.make_ptr(self.io_dtype, 0, cute.AddressSpace.smem), cute.dice(self.kv_mma_tiler, (1, 1, None)), @@ -4372,10 +4735,17 @@ def make_tmem_load_and_partition(self, copy_atom_t2r, tmem_tensor, tmem_tile_coo def tmem_store_and_partition_acc(self, local_tidx, tCtAcc): dtype = tCtAcc.element_type - copy_atom_r2t = cute.make_copy_atom( - tcgen05.St32x32bOp(tcgen05.Repetition(32), tcgen05.Unpack.NONE), - dtype, - ) + if cutlass.const_expr(self.split_value_tiles): + # Inverse of the split state's 16dp Ld16x256b x16 mapping. + copy_atom_r2t = cute.make_copy_atom( + tcgen05.St16x256bOp(tcgen05.Repetition(16), tcgen05.Unpack.NONE), + dtype, + ) + else: + copy_atom_r2t = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32), tcgen05.Unpack.NONE), + dtype, + ) tiled_r2t = tcgen05.make_tmem_copy(copy_atom_r2t, tCtAcc) thr_r2t = tiled_r2t.get_slice(local_tidx) @@ -4616,22 +4986,25 @@ def epilog_tmem_copy_and_partition( """ # Make tiledCopy for tensor memory load epitile = mma_tiler[:2] - assert epitile[0] == 128 + assert epitile[0] == self.value_tile_size # TODO: 32dp ease DEBUGGING - copy_atom_t2r = cute.make_copy_atom( - tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), - self.acc_dtype, - ) - # copy_atom_t2r = sm100_utils.get_tmem_load_op( - # mma_tiler, - # # self.o_layout, - # # TODO - # utils.LayoutEnum.ROW_MAJOR, - # self.io_dtype, - # self.acc_dtype, - # epitile, - # use_2cta_instrs, - # ) + if cutlass.const_expr(self.split_value_tiles): + # M=64 has only 16 TMEM datapaths per warp. Let CUTLASS select the + # matching 16dp load atom (currently Ld16x256b) instead of forcing + # the 32dp atom used by the original M=128 epilogue. + copy_atom_t2r = sm100_utils.get_tmem_load_op( + mma_tiler, + utils.LayoutEnum.ROW_MAJOR, + self.io_dtype, + self.acc_dtype, + epitile, + use_2cta_instrs, + ) + else: + copy_atom_t2r = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32), tcgen05.Pack.NONE), + self.acc_dtype, + ) # (EPI_TILE_M, EPI_TILE_N, 1, 1, STAGE) tAcc_epi = cute.flat_divide( # ((EPI_TILE_M, EPI_TILE_N), EPI_M, EPI_N, STAGE) @@ -5610,6 +5983,191 @@ def tma_partition_for_mma_operand( # =========== # Utility functions for Ampere-style mma.sync, used for subchunk computation + @cute.jit + def produce_scalar_gate_chunk( + self, + g: cute.Tensor, + sG_tma: cute.Tensor, + sG_factor_a: cute.Tensor, + sG_factor_b: cute.Tensor, + sG_factor_base: cute.Tensor, + sG_main_q: cute.Tensor, + sG_main_k: cute.Tensor, + chunk_start, + seq_len, + tok_offset, + hidx, + data_bidx, + lane_idx, + stage_idx, + ): + """Produce one scalar-G stage with a single 32-thread warp.""" + for lane_slot in cutlass.range_constexpr(2): + token_in_chunk = lane_idx + lane_slot * self.threads_per_warp + token_in_seq = chunk_start + token_in_chunk + if token_in_seq < seq_len: + token_in_data = tok_offset + token_in_seq + sG_tma[token_in_chunk, 0, stage_idx] = g[ + token_in_data, + 0, + (hidx, data_bidx), + ] + else: + sG_tma[token_in_chunk, 0, stage_idx] = cutlass.Float32(0.0) + cute.arch.sync_warp() + + if cutlass.const_expr(self.fuse_scalar_cumsum): + # Inclusive chunk-local cumsum in log2 space. The producer warp + # owns exactly two 32-token halves. + g_lo = sG_tma[lane_idx, 0, stage_idx] * Constant.RCP_LN2 + g_hi = ( + sG_tma[lane_idx + self.threads_per_warp, 0, stage_idx] + * Constant.RCP_LN2 + ) + for scan_step in cutlass.range_constexpr(5): + scan_offset = 1 << scan_step + add_lo = cute.arch.shuffle_sync_up( + g_lo, + scan_offset, + mask_and_clamp=0, + ) + add_hi = cute.arch.shuffle_sync_up( + g_hi, + scan_offset, + mask_and_clamp=0, + ) + if lane_idx >= scan_offset: + g_lo += add_lo + g_hi += add_hi + g_lo_total = cute.arch.shuffle_sync(g_lo, 31) + g_hi += g_lo_total + sG_tma[lane_idx, 0, stage_idx] = g_lo + sG_tma[lane_idx + self.threads_per_warp, 0, stage_idx] = g_hi + cute.arch.sync_warp() + + # Factor pairwise gates as: + # A[t] = exp2(g[t]-g[key_subchunk_base]) + # Base[q,k] = exp2(g[q_base]-g[k_base]) + # B[q,t] = Base[q,key_subchunk(t)] / A[t] + valid_len_g = seq_len - chunk_start + if valid_len_g > Constant.C: + valid_len_g = Constant.C + g_last_scalar = sG_tma[valid_len_g - 1, 0, stage_idx] + for lane_slot in cutlass.range_constexpr(2): + token_in_chunk = lane_idx + lane_slot * self.threads_per_warp + if token_in_chunk < valid_len_g: + g_i = sG_tma[token_in_chunk, 0, stage_idx] + own_base = (token_in_chunk // Constant.SC) * Constant.SC + g_own_base = sG_tma[own_base, 0, stage_idx] + sG_factor_a[token_in_chunk, stage_idx] = cute.exp2( + g_i - g_own_base, + fastmath=self.use_fast_math, + ) + sG_main_q[token_in_chunk, stage_idx] = cute.exp2( + g_i, + fastmath=self.use_fast_math, + ) + sG_main_k[token_in_chunk, stage_idx] = cute.exp2( + g_last_scalar - g_i, + fastmath=self.use_fast_math, + ) + else: + sG_factor_a[token_in_chunk, stage_idx] = cutlass.Float32(0.0) + sG_main_q[token_in_chunk, stage_idx] = cutlass.Float32(0.0) + sG_main_k[token_in_chunk, stage_idx] = cutlass.Float32(0.0) + + num_subchunks = Constant.C // Constant.SC + if lane_idx < num_subchunks * num_subchunks: + query_subchunk = lane_idx // num_subchunks + key_subchunk = lane_idx % num_subchunks + query_base = query_subchunk * Constant.SC + key_base = key_subchunk * Constant.SC + if ( + key_subchunk <= query_subchunk + and query_base < valid_len_g + and key_base < valid_len_g + ): + sG_factor_base[ + query_subchunk, + key_subchunk, + stage_idx, + ] = cute.exp2( + sG_tma[query_base, 0, stage_idx] + - sG_tma[key_base, 0, stage_idx], + fastmath=self.use_fast_math, + ) + else: + sG_factor_base[ + query_subchunk, + key_subchunk, + stage_idx, + ] = cutlass.Float32(0.0) + cute.arch.sync_warp() + + for lane_slot in cutlass.range_constexpr(2): + token_in_chunk = lane_idx + lane_slot * self.threads_per_warp + if token_in_chunk < valid_len_g: + key_subchunk = token_in_chunk // Constant.SC + inv_a = cutlass.Float32(1.0) / sG_factor_a[ + token_in_chunk, + stage_idx, + ] + for query_subchunk in cutlass.range_constexpr(num_subchunks): + if key_subchunk <= query_subchunk: + sG_factor_b[ + query_subchunk, + token_in_chunk, + stage_idx, + ] = ( + sG_factor_base[ + query_subchunk, + key_subchunk, + stage_idx, + ] + * inv_a + ) + else: + sG_factor_b[ + query_subchunk, + token_in_chunk, + stage_idx, + ] = cutlass.Float32(0.0) + else: + for query_subchunk in cutlass.range_constexpr(num_subchunks): + sG_factor_b[ + query_subchunk, + token_in_chunk, + stage_idx, + ] = cutlass.Float32(0.0) + + @cute.jit + def apply_scalar_gate_to_subchunk_acc( + self, + qk_acc: cute.Tensor, + kk_acc: cute.Tensor, + acc_coords: cute.Tensor, + scalar_g: cute.Tensor, + row_subchunk: cutlass.Constexpr[int], + col_subchunk: cutlass.Constexpr[int], + ): + """Apply exp2(g[row]-g[col]) to FP32 QK/KK accumulators. + + The vector-gate implementation factors this term into separately + scaled BF16 Q and K operands. For a native scalar gate, applying the + exact token-pair factor to the FP32 accumulator is both cheaper and + avoids relying on operand-fragment coordinates to recover token ids. + """ + row_base = row_subchunk * Constant.SC + col_base = col_subchunk * Constant.SC + for i in cutlass.range_constexpr(cute.size(acc_coords)): + row, col = acc_coords[i] + gate = cute.exp2( + scalar_g[row_base + row] - scalar_g[col_base + col], + fastmath=self.use_fast_math, + ) + qk_acc[i] *= gate + kk_acc[i] *= gate + @cute.jit def mma_sync_partition_c( self, tiled_mma: cute.atom.TiledMma, tile_shape_mnk: cute.Shape, zero_fill: cutlass.Constexpr[bool] = True @@ -5630,31 +6188,42 @@ def s2r_compute_subchunk_operand_A( q_k_tiled_copy: cute.atom.TiledCopy, q_k_thr_copy: cute.atom.ThrCopy, tv_layout_mma_A: cute.Layout, + scalar_coords: cute.Tensor, layout_g_first: cute.Layout, # for make g_first tensor sG_slice: cute.Tensor, + sG_scalar: cute.Tensor, + sG_factor_a: cute.Tensor, sQ_slice: cute.Tensor, sK_slice: cute.Tensor, ): - # S2R g, g_first - sG = sG_slice[None, None, subchunk_idx, (0, nk)] - tQKsG = g_thr_copy.partition_S(sG) tQKrG = cute.make_fragment_like(tv_layout_mma_A, dtype=self.g_dtype) - tQKrG_cv = g_thr_copy.retile(tQKrG) - cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) - - # TODO: do register shuffle g to get g_first, reduce smem load - sG_first = cute.make_tensor(sG.iterator, layout=layout_g_first) - tQKsGfirst = g_thr_copy.partition_S(sG_first) - tQKrGfirst = cute.make_fragment_like(tv_layout_mma_A, dtype=self.g_dtype) - tQKrGfirst_cv = g_thr_copy.retile(tQKrGfirst) - cute.copy(g_tiled_copy, tQKsGfirst, tQKrGfirst_cv) + if cutlass.const_expr(self.scalar_gate): + # Coordinates come from the exact MMA-A partition, matching the + # flattened register fragment one-for-one. + for i in cutlass.range_constexpr(cute.size(scalar_coords)): + index_q, _ = scalar_coords[i] + tQKrG[i] = sG_factor_a[subchunk_idx * Constant.SC + index_q] + else: + # S2R g, g_first + sG = sG_slice[None, None, subchunk_idx, (0, nk)] + tQKsG = g_thr_copy.partition_S(sG) + tQKrG_cv = g_thr_copy.retile(tQKrG) + cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) + + # TODO: do register shuffle g to get g_first, reduce smem load + sG_first = cute.make_tensor(sG.iterator, layout=layout_g_first) + tQKsGfirst = g_thr_copy.partition_S(sG_first) + tQKrGfirst = cute.make_fragment_like(tv_layout_mma_A, dtype=self.g_dtype) + tQKrGfirst_cv = g_thr_copy.retile(tQKrGfirst) + cute.copy(g_tiled_copy, tQKsGfirst, tQKrGfirst_cv) + + # gqn = exp2(g - g_first[None, :]), reuse g + g_val = tQKrG.load() + g_first_val = tQKrGfirst.load() + g_val = cute.exp2(g_val - g_first_val, fastmath=self.use_fast_math) + tQKrG.store(g_val) - # gqn = exp2(g - g_first[None, :]), reuse g g_val = tQKrG.load() - g_first_val = tQKrGfirst.load() - g_val = cute.exp2(g_val - g_first_val, fastmath=self.use_fast_math) - tQKrG.store(g_val) - # S2R q, k sQ = sQ_slice[None, None, subchunk_idx, (0, nk)] sK = sK_slice[None, None, subchunk_idx, (0, nk)] @@ -5689,23 +6258,36 @@ def s2r_compute_subchunk_operand_B( kt_tiled_copy: cute.atom.TiledCopy, kt_thr_copy: cute.atom.ThrCopy, tv_layout_mma_B: cute.Layout, + scalar_coords: cute.Tensor, sG_slice: cute.Tensor, + sG_scalar: cute.Tensor, + sG_factor_b: cute.Tensor, + query_subchunk_idx: cutlass.Constexpr[int], sK_slice: cute.Tensor, rG_first: cute.Tensor, ): - # S2R g - sG = sG_slice[None, None, subchunk_idx, (0, nk)] - tQKsG = g_thr_copy.partition_S(sG) tQKrG = cute.make_fragment_like(tv_layout_mma_B, dtype=self.g_dtype) - tQKrG_cv = g_thr_copy.retile(tQKrG) - cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) + if cutlass.const_expr(self.scalar_gate): + for i in cutlass.range_constexpr(cute.size(scalar_coords)): + index_k, _ = scalar_coords[i] + tQKrG[i] = sG_factor_b[ + query_subchunk_idx, + subchunk_idx * Constant.SC + index_k, + ] + else: + # S2R g + sG = sG_slice[None, None, subchunk_idx, (0, nk)] + tQKsG = g_thr_copy.partition_S(sG) + tQKrG_cv = g_thr_copy.retile(tQKrG) + cute.copy(g_tiled_copy, tQKsG, tQKrG_cv) + + # compute gktn = exp2(g_first - g), reuse g + g_val = tQKrG.load() + g_first_val = rG_first.load() + g_val = cute.exp2(g_first_val - g_val, fastmath=self.use_fast_math) + tQKrG.store(g_val) - # compute gktn = exp2(g_first - g), reuse g g_val = tQKrG.load() - g_first_val = rG_first.load() - g_val = cute.exp2(g_first_val - g_val, fastmath=self.use_fast_math) - tQKrG.store(g_val) - # S2R k sK = sK_slice[None, None, subchunk_idx, (0, nk)] tQKrKt = cute.make_fragment_like(tv_layout_mma_B, dtype=self.k_dtype) @@ -5791,7 +6373,7 @@ def make_s2r_partitions_v( # num_bits_per_copy=dtype.width * 8, num_bits_per_copy=dtype.width * 1, ) - num_elements_per_thread = Constant.C + num_elements_per_thread = Constant.C // self.num_value_tiles num_threads_per_row = shape_x[1] // num_elements_per_thread # NOTE: Assume 128 cuda core threads num_threads_per_col = 128 // num_threads_per_row diff --git a/cula/ops/qwen35_fused_kda_prefill.py b/cula/ops/qwen35_fused_kda_prefill.py index 6691fdd8..39f32024 100644 --- a/cula/ops/qwen35_fused_kda_prefill.py +++ b/cula/ops/qwen35_fused_kda_prefill.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Qwen3.5 adapter for the generic fused KDA prefill core.""" +"""Qwen3.5 adapter for the native-GVA fully-fused CuTe prefill core.""" from __future__ import annotations @@ -52,12 +52,19 @@ def _validate_inputs( cu_seqlens: torch.Tensor | None, ) -> tuple[int, int, int, int, torch.Tensor, torch.Tensor]: if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: - raise ValueError(f"q/k/v must be 4D [B,T,HV,D], got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") - if q.shape != k.shape or q.shape != v.shape: - raise ValueError(f"q/k/v must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") - B, T, HV, K = q.shape - if K != 128: - raise ValueError(f"Qwen3.5 fused prefill expects head dim 128, got {K}") + raise ValueError( + f"q/k/v must be 4D [B,T,H,D], got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}" + ) + if q.shape != k.shape: + raise ValueError(f"q and k must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)}") + B, T, H, K = q.shape + if v.shape[:2] != (B, T): + raise ValueError(f"v must match q/k batch and token dimensions, got q={tuple(q.shape)} v={tuple(v.shape)}") + HV, V = v.shape[2:] + if K != 128 or V != 128: + raise ValueError(f"Qwen3.5 fused prefill expects K=V=128, got K={K} V={V}") + if HV % H != 0: + raise ValueError(f"Qwen3.5 GVA expects HV to be divisible by H, got H={H} HV={HV}") if a.ndim == 2: a = a.unsqueeze(0) if b.ndim == 2: @@ -72,8 +79,8 @@ def _validate_inputs( if cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: raise ValueError(f"cu_seqlens must be 1D int32, got {tuple(cu_seqlens.shape)} {cu_seqlens.dtype}") state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 - if initial_state is not None and initial_state.shape != (state_count, HV, K, K): - raise ValueError(f"initial_state must be [{state_count},{HV},128,128], got {tuple(initial_state.shape)}") + if initial_state is not None and initial_state.shape != (state_count, HV, K, V): + raise ValueError(f"initial_state must be [{state_count},{HV},{K},{V}], got {tuple(initial_state.shape)}") return B, T, HV, K, a, b @@ -90,25 +97,37 @@ def qwen35_fused_kda_prefill( cu_seqlens: torch.Tensor | None = None, output_final_state: bool = True, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Run Qwen3.5 scalar-gated KDA prefill through the fused CuTe KDA core. + """Run Qwen3.5 scalar-gated KDA prefill through the fully-fused CuTe core. + + Qwen3.5 uses native grouped value attention: Q/K have H heads while V and + the scalar GDN gate have HV heads (globally H=16 and HV=48). The scalar + gate is passed to the CuTe specialization without materializing a D=128 + broadcast; Q/K likewise remain in their native, non-repeated layout. - Qwen uses a scalar gate per token/head. The generic KDA fused core expects - a vector gate, so this adapter broadcasts the scalar log-gate over D=128. State is exposed in Qwen layout [N, HV, K, V]. The fused core consumes the - transposed initial-state layout, but returns final state in Qwen layout. + transposed initial state and returns final state in Qwen layout. """ if not q.is_cuda: raise RuntimeError("qwen35_fused_kda_prefill requires CUDA tensors.") B, T, HV, K, a, b = _validate_inputs(q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens) + fused_kda_prefill = _resolve_fused_kda_prefill(q.device) log_gate_scalar = -torch.exp(A_log.float()).view(1, 1, HV, 1) * F.softplus( a.float().unsqueeze(-1) + dt_bias.float().view(1, 1, HV, 1) ) - log_gate = log_gate_scalar.expand(B, T, HV, K).contiguous() + log_gate = log_gate_scalar.squeeze(-1).contiguous() beta = torch.sigmoid(b.float()).contiguous() + # A single [0, T] sequence is an ordinary equal-length prefill, not a + # variable-length batch. Keep the fast non-varlen SM100 launch in this + # common Qwen inference case; the varlen path carries extra indirection + # and workspace overhead. + kernel_cu_seqlens = cu_seqlens + if cu_seqlens is not None and q.shape[0] == 1 and cu_seqlens.numel() == 2: + kernel_cu_seqlens = None + initial_state_vk = None if initial_state is not None: initial_state_vk = initial_state.float().transpose(-1, -2).contiguous() @@ -124,9 +143,10 @@ def qwen35_fused_kda_prefill( output_final_state=output_final_state, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=False, - safe_gate=False, - lower_bound=None, - cu_seqlens=cu_seqlens, + cu_seqlens=kernel_cu_seqlens, + safe_gate=True, + lower_bound=-5.0, + scalar_gate=True, ) final_state = None if final_state_vk is None else final_state_vk.contiguous() return out, final_state From 3bdb09a8cbf9f56fd9ef85ec61b951b26e6a5f43 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sun, 2 Aug 2026 15:38:55 +0000 Subject: [PATCH 21/35] test(qwen35): cover native GVA fused prefill --- tests/test_qwen35_prefill.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/test_qwen35_prefill.py b/tests/test_qwen35_prefill.py index 445c9369..1bca86ff 100644 --- a/tests/test_qwen35_prefill.py +++ b/tests/test_qwen35_prefill.py @@ -280,7 +280,8 @@ def test_qwen35_chunk_qk_prefill_sm90_supports_local_tp_shards(local_v_heads: in torch.testing.assert_close(out, ref, atol=2e-1, rtol=2e-2) -def test_qwen35_fused_kda_prefill_matches_reference(): +@pytest.mark.parametrize("T", [64, 128]) +def test_qwen35_fused_kda_prefill_matches_reference(T: int): if not torch.cuda.is_available(): import pytest @@ -292,10 +293,10 @@ def test_qwen35_fused_kda_prefill_matches_reference(): torch.manual_seed(12) device = torch.device("cuda") - B, T, HV, K = 1, 64, 48, 128 - q = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + B, H, HV, K = 1, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) k = torch.randn_like(q) - v = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) b = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) A_log = -torch.rand(HV, device=device, dtype=torch.float32) @@ -303,8 +304,8 @@ def test_qwen35_fused_kda_prefill_matches_reference(): initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 out_ref, state_ref = qwen35_scalar_kda_prefill( - q, - k, + q.repeat_interleave(HV // H, dim=2), + k.repeat_interleave(HV // H, dim=2), v, a, b, From cf382eb0614987d187795ee2a5d9b48506620d1f Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sun, 2 Aug 2026 15:39:06 +0000 Subject: [PATCH 22/35] bench(qwen35): compare actual GDN prefill paths --- benchmarks/bench_qwen35_prefill.py | 252 +++++++++++++++++++++++++++++ 1 file changed, 252 insertions(+) create mode 100644 benchmarks/bench_qwen35_prefill.py diff --git a/benchmarks/bench_qwen35_prefill.py b/benchmarks/bench_qwen35_prefill.py new file mode 100644 index 00000000..34cb989f --- /dev/null +++ b/benchmarks/bench_qwen35_prefill.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +"""Benchmark cuLA native-GVA Qwen3.5 prefill against SGLang's inference path. + +Only config.json is read. Model weights are not loaded: tensors are generated +from the Qwen3.5 linear-attention shapes and dtype declared by the config. +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import statistics +import sys +from collections.abc import Callable + +import torch + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +from cula.ops.qwen35_fused_kda_prefill import qwen35_fused_kda_prefill + + +def load_qwen35_shape(config_path: pathlib.Path, tp_size: int) -> dict[str, object]: + with config_path.open(encoding="utf-8") as f: + root = json.load(f) + config = root.get("text_config", root) + + required = ( + "linear_num_key_heads", + "linear_num_value_heads", + "linear_key_head_dim", + "linear_value_head_dim", + ) + missing = [key for key in required if key not in config] + if missing: + raise ValueError(f"{config_path} is missing Qwen3.5 fields: {missing}") + + global_h = int(config["linear_num_key_heads"]) + global_hv = int(config["linear_num_value_heads"]) + if global_h % tp_size or global_hv % tp_size: + raise ValueError(f"TP={tp_size} must divide H={global_h} and HV={global_hv}") + + h = global_h // tp_size + hv = global_hv // tp_size + k = int(config["linear_key_head_dim"]) + v = int(config["linear_value_head_dim"]) + if hv % h: + raise ValueError(f"Qwen3.5 GVA requires local HV % H == 0, got H={h} HV={hv}") + if k != 128 or v != 128: + raise ValueError(f"cuLA native-GVA prefill currently requires K=V=128, got K={k} V={v}") + + dtype_name = str( + config.get("torch_dtype", config.get("dtype", root.get("torch_dtype", root.get("dtype", "bfloat16")))) + ).lower() + if dtype_name not in ("bfloat16", "bf16", "torch.bfloat16"): + raise ValueError(f"This benchmark expects Qwen3.5 bf16 activations, got torch_dtype={dtype_name}") + + return { + "model_type": config.get("model_type", root.get("model_type", "unknown")), + "global_h": global_h, + "global_hv": global_hv, + "h": h, + "hv": hv, + "k": k, + "v": v, + "dtype": torch.bfloat16, + } + + +def load_sglang(sglang_path: pathlib.Path | None): + if sglang_path is not None: + for candidate in (sglang_path, sglang_path / "python"): + if candidate.exists(): + sys.path.insert(0, str(candidate)) + + from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating + from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel + + return fused_gdn_gating, TritonGDNKernel() + + +def benchmark_cuda( + fn: Callable[[], object], + warmup: int, + rep: int, + setup: Callable[[], None] | None = None, +) -> float: + for _ in range(warmup): + if setup is not None: + setup() + fn() + torch.cuda.synchronize() + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] + for start, end in zip(starts, ends, strict=True): + if setup is not None: + setup() + start.record() + fn() + end.record() + torch.cuda.synchronize() + + samples = sorted(start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)) + if len(samples) < 4: + return statistics.mean(samples) + return statistics.mean(samples[len(samples) // 4 : 3 * len(samples) // 4]) + + +def relative_rms(ref: torch.Tensor, out: torch.Tensor) -> float: + ref_f = ref.float() + diff_rms = (ref_f - out.float()).square().mean().sqrt() + return (diff_rms / ref_f.square().mean().sqrt().clamp_min(1.0e-8)).item() + + +def make_inputs( + batch: int, + seq_len: int, + shape: dict[str, object], + device: torch.device, + seed: int, + random_initial_state: bool, +) -> dict[str, torch.Tensor]: + torch.manual_seed(seed) + total = batch * seq_len + h, hv, k, v = (int(shape[name]) for name in ("h", "hv", "k", "v")) + dtype = shape["dtype"] + + state = torch.zeros(batch, hv, k, v, device=device, dtype=torch.float32) + if random_initial_state: + state.normal_(mean=0.0, std=0.01) + + return { + "q": torch.randn(1, total, h, k, device=device, dtype=dtype), + "k": torch.randn(1, total, h, k, device=device, dtype=dtype), + "v": torch.randn(1, total, hv, v, device=device, dtype=dtype), + "a": torch.randn(total, hv, device=device, dtype=dtype), + "b": torch.randn(total, hv, device=device, dtype=dtype), + "A_log": -torch.rand(hv, device=device, dtype=torch.float32), + "dt_bias": torch.randn(hv, device=device, dtype=torch.float32) * 0.1, + "state_kv": state, + "state_vk": state.transpose(-1, -2).contiguous(), + "cu_seqlens": torch.arange(0, total + 1, seq_len, device=device, dtype=torch.int32), + "cache_indices": torch.arange(batch, device=device, dtype=torch.int32), + } + + +@torch.inference_mode() +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--sglang-path", type=pathlib.Path, default=pathlib.Path("/sgl-workspace/sglang")) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-lens", type=int, nargs="+", default=(128, 256, 512, 1024, 2048, 4096)) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rep", type=int, default=30) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--random-initial-state", action="store_true") + parser.add_argument("--skip-accuracy", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + + shape = load_qwen35_shape(args.config_json, args.tp_size) + fused_gdn_gating, sglang_kernel = load_sglang(args.sglang_path) + device = torch.device("cuda") + + print("Qwen3.5 GDN prefill: cuLA native GVA vs SGLang Triton inference path") + print(f"config={args.config_json} model_type={shape['model_type']}") + print( + f"device={torch.cuda.get_device_name(device)} TP={args.tp_size} " + f"global H/HV={shape['global_h']}/{shape['global_hv']} " + f"local H/HV={shape['h']}/{shape['hv']} K/V={shape['k']}/{shape['v']}" + ) + print(f"batch={args.batch} warmup={args.warmup} rep={args.rep}") + print() + print(f"{'T/seq':>8} {'tokens':>8} {'SGLang ms':>11} {'cuLA ms':>10} {'speedup':>9} {'out rrms':>11} {'state rrms':>12}") + print("-" * 79) + + for seq_len in args.seq_lens: + x = make_inputs( + args.batch, + seq_len, + shape, + device, + args.seed, + args.random_initial_state, + ) + state_sglang = torch.empty_like(x["state_vk"]) + + def setup_sglang(): + state_sglang.copy_(x["state_vk"]) + + def run_cula(): + return qwen35_fused_kda_prefill( + x["q"], + x["k"], + x["v"], + x["a"], + x["b"], + x["A_log"], + x["dt_bias"], + initial_state=x["state_kv"], + cu_seqlens=x["cu_seqlens"], + output_final_state=True, + ) + + def run_sglang(): + g, beta = fused_gdn_gating( + x["A_log"], + x["a"], + x["b"], + x["dt_bias"], + ) + return sglang_kernel.extend( + x["q"], + x["k"], + x["v"], + g, + beta, + ssm_states=state_sglang, + cache_indices=x["cache_indices"], + query_start_loc=x["cu_seqlens"], + ) + + rrms = float("nan") + state_rrms = float("nan") + if not args.skip_accuracy: + out_cula, state_cula = run_cula() + setup_sglang() + out_sglang = run_sglang()[0] + torch.cuda.synchronize() + rrms = relative_rms(out_sglang, out_cula) + state_rrms = relative_rms(state_sglang, state_cula.transpose(-1, -2)) + + sglang_ms = benchmark_cuda(run_sglang, args.warmup, args.rep, setup=setup_sglang) + cula_ms = benchmark_cuda(run_cula, args.warmup, args.rep) + total = args.batch * seq_len + print( + f"{seq_len:8d} {total:8d} {sglang_ms:11.4f} {cula_ms:10.4f} " + f"{sglang_ms / cula_ms:8.3f}x {rrms:11.3e} {state_rrms:12.3e}" + ) + del x + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() From 020172b7b4770d259fb1dcedbbf489961f3c3209 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sun, 2 Aug 2026 15:52:49 +0000 Subject: [PATCH 23/35] bench(qwen35): compare native GVA decode inputs --- benchmarks/bench_qwen35_decode.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/benchmarks/bench_qwen35_decode.py b/benchmarks/bench_qwen35_decode.py index dbf4cbc2..40e934cb 100755 --- a/benchmarks/bench_qwen35_decode.py +++ b/benchmarks/bench_qwen35_decode.py @@ -20,7 +20,6 @@ from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule import cula.cudac as cula_cuda -from cula.ops.qwen35_layout_decode import qwen35_layout_decode_reference from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as GLOBAL_CONFIG from cula.qwen35.common import Qwen35LinearAttentionConfig @@ -104,10 +103,15 @@ def run_case(tokens: int, args, device: torch.device) -> dict[str, float | int]: seed=args.seed, device=device, ) - q, k, v, a_fla, b_fla = qwen35_layout_decode_reference(mixed, a, b, config=config) + qk_width = config.num_k_heads * config.head_k_dim + q = mixed[:, :qk_width].view(tokens, config.num_k_heads, config.head_k_dim).contiguous() + k = mixed[:, qk_width : 2 * qk_width].view( + tokens, config.num_k_heads, config.head_k_dim + ).contiguous() + v = mixed[:, 2 * qk_width :].view(tokens, config.num_v_heads, config.head_v_dim).contiguous() q, k, v = q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1) - gate = a_fla.unsqueeze(1) - beta = torch.sigmoid(b_fla.float()).unsqueeze(1) + gate = a.contiguous().unsqueeze(1) + beta = torch.sigmoid(b.float()).unsqueeze(1) state_cula = torch.empty_like(state) out_cula = torch.empty_like(v.squeeze(1)) From 39e2602053d57223896ed5c735e107b3576a3234 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sun, 2 Aug 2026 16:02:20 +0000 Subject: [PATCH 24/35] bench(qwen35): compare SGLang packed decode --- benchmarks/bench_qwen35_decode.py | 270 ++++++++++++++---------------- 1 file changed, 129 insertions(+), 141 deletions(-) diff --git a/benchmarks/bench_qwen35_decode.py b/benchmarks/bench_qwen35_decode.py index 40e934cb..bfbdcd07 100755 --- a/benchmarks/bench_qwen35_decode.py +++ b/benchmarks/bench_qwen35_decode.py @@ -1,211 +1,199 @@ #!/usr/bin/env python3 -"""Benchmark the fused cuLA Qwen3.5 decode kernel against upstream FLA. +"""Benchmark actual cuLA Qwen GDN decode against SGLang's packed inference path. -The upstream FLA recurrent operator receives pre-laid-out Q/K/V tensors. cuLA -receives the packed Qwen3.5 ``mixed_qkv_conv`` tensor and performs layout plus -the recurrent update in one kernel. State reset is outside both timing windows. +Only config.json is read. State reset is outside both CUDA event windows. """ from __future__ import annotations import argparse import csv +import json import pathlib import statistics import sys - -sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +from collections.abc import Callable import torch -from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + +ROOT = pathlib.Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) import cula.cudac as cula_cuda -from cula.qwen35.common import DEFAULT_QWEN35_LINEAR_ATTN_CONFIG as GLOBAL_CONFIG -from cula.qwen35.common import Qwen35LinearAttentionConfig - - -def local_config_from_tp_size(tp_size: int) -> Qwen35LinearAttentionConfig: - return Qwen35LinearAttentionConfig( - hidden_size=GLOBAL_CONFIG.hidden_size // tp_size, - conv_kernel_size=GLOBAL_CONFIG.conv_kernel_size, - num_k_heads=GLOBAL_CONFIG.num_k_heads // tp_size, - num_v_heads=GLOBAL_CONFIG.num_v_heads // tp_size, - head_k_dim=GLOBAL_CONFIG.head_k_dim, - head_v_dim=GLOBAL_CONFIG.head_v_dim, - qkv_dtype=GLOBAL_CONFIG.qkv_dtype, - state_dtype=GLOBAL_CONFIG.state_dtype, - ) -def benchmark_cuda(fn, *, setup=None, warmup: int, rep: int) -> float: +def load_shape(config_path: pathlib.Path, tp_size: int) -> dict[str, int | str]: + with config_path.open(encoding="utf-8") as f: + root = json.load(f) + config = root.get("text_config", root) + global_h = int(config["linear_num_key_heads"]) + global_hv = int(config["linear_num_value_heads"]) + if global_h % tp_size or global_hv % tp_size: + raise ValueError(f"TP={tp_size} must divide H={global_h} and HV={global_hv}") + h, hv = global_h // tp_size, global_hv // tp_size + k = int(config["linear_key_head_dim"]) + v = int(config["linear_value_head_dim"]) + if hv % h or k != 128 or v != 128: + raise ValueError(f"unsupported local GVA shape H={h} HV={hv} K={k} V={v}") + return { + "model": config_path.parent.name, + "global_h": global_h, + "global_hv": global_hv, + "h": h, + "hv": hv, + "k": k, + "v": v, + } + + +def load_sglang(sglang_path: pathlib.Path): + for candidate in (sglang_path, sglang_path / "python"): + if candidate.exists(): + sys.path.insert(0, str(candidate)) + from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel + + kernel = TritonGDNKernel() + if not kernel.supports_packed_decode: + raise RuntimeError("SGLang Triton packed GDN decode is unavailable") + return kernel + + +def benchmark_cuda( + fn: Callable[[], object], + *, + setup: Callable[[], None], + warmup: int, + rep: int, +) -> float: for _ in range(warmup): - if setup is not None: - setup() + setup() fn() torch.cuda.synchronize() starts = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] ends = [torch.cuda.Event(enable_timing=True) for _ in range(rep)] - for i in range(rep): - if setup is not None: - setup() - starts[i].record() + for start, end in zip(starts, ends, strict=True): + setup() + start.record() fn() - ends[i].record() + end.record() torch.cuda.synchronize() + samples = sorted(start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)) + if len(samples) < 4: + return statistics.mean(samples) + return statistics.mean(samples[len(samples) // 4 : 3 * len(samples) // 4]) + - times = sorted(start.elapsed_time(end) for start, end in zip(starts, ends)) - lo, hi = len(times) // 4, 3 * len(times) // 4 - return statistics.mean(times[lo:hi] or times) +def relative_rms(reference: torch.Tensor, actual: torch.Tensor) -> float: + ref = reference.float() + diff = ref - actual.float() + return (diff.square().mean().sqrt() / ref.square().mean().sqrt().clamp_min(1e-8)).item() -def error_stats(reference: torch.Tensor, actual: torch.Tensor) -> tuple[float, float]: - reference = reference.float() - actual = actual.float() - diff = (reference - actual).abs() - rel_rms = diff.square().mean().sqrt() / reference.square().mean().sqrt().clamp_min(1e-8) - return rel_rms.item(), diff.max().item() +def make_inputs(tokens: int, shape: dict[str, int | str], seed: int) -> dict[str, torch.Tensor]: + torch.manual_seed(seed) + device = torch.device("cuda") + h, hv, k, v = (int(shape[name]) for name in ("h", "hv", "k", "v")) + conv_dim = 2 * h * k + hv * v + state_kv = torch.randn(tokens, hv, k, v, device=device, dtype=torch.float32) * 0.01 + return { + "mixed_qkv": torch.randn(tokens, conv_dim, device=device, dtype=torch.bfloat16), + "a": torch.randn(tokens, hv, device=device, dtype=torch.bfloat16), + "b": torch.randn(tokens, hv, device=device, dtype=torch.bfloat16), + "A_log": -torch.rand(hv, device=device, dtype=torch.float32), + "dt_bias": torch.randn(hv, device=device, dtype=torch.float32) * 0.1, + "state_kv": state_kv, + "state_vk": state_kv.transpose(-1, -2).contiguous(), + "indices": torch.arange(tokens, device=device, dtype=torch.int32), + } -def make_inputs(tokens: int, *, tp_size: int, seed: int, device: torch.device): - config = local_config_from_tp_size(tp_size) - generator = torch.Generator(device=device).manual_seed(seed) - hv, k_dim, v_dim = config.num_v_heads, config.head_k_dim, config.head_v_dim +@torch.inference_mode() +def run_case(tokens: int, shape, sglang_kernel, args) -> dict[str, float | int]: + x = make_inputs(tokens, shape, args.seed) + hv, k, v = (int(shape[name]) for name in ("hv", "k", "v")) + state_cula = torch.empty_like(x["state_kv"]) + state_sglang = torch.empty_like(x["state_vk"]) + out_cula = torch.empty(tokens, hv, v, device="cuda", dtype=torch.bfloat16) - mixed_qkv = torch.randn( - tokens, - config.conv_dim, - generator=generator, - device=device, - dtype=config.qkv_dtype, - ) - a = torch.randn(tokens, hv, generator=generator, device=device, dtype=config.qkv_dtype) - b = torch.randn(tokens, hv, generator=generator, device=device, dtype=config.qkv_dtype) - A_log = -torch.rand(hv, generator=generator, device=device, dtype=torch.float32) - dt_bias = torch.randn(hv, generator=generator, device=device, dtype=torch.float32) * 0.1 - state = torch.randn( - tokens, - hv, - k_dim, - v_dim, - generator=generator, - device=device, - dtype=torch.float32, - ) * 0.01 - indices = torch.arange(tokens, device=device, dtype=torch.int32) - return config, mixed_qkv, a, b, A_log, dt_bias, state, indices - - -def run_case(tokens: int, args, device: torch.device) -> dict[str, float | int]: - config, mixed, a, b, A_log, dt_bias, state, indices = make_inputs( - tokens, - tp_size=args.tp_size, - seed=args.seed, - device=device, - ) - qk_width = config.num_k_heads * config.head_k_dim - q = mixed[:, :qk_width].view(tokens, config.num_k_heads, config.head_k_dim).contiguous() - k = mixed[:, qk_width : 2 * qk_width].view( - tokens, config.num_k_heads, config.head_k_dim - ).contiguous() - v = mixed[:, 2 * qk_width :].view(tokens, config.num_v_heads, config.head_v_dim).contiguous() - q, k, v = q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1) - gate = a.contiguous().unsqueeze(1) - beta = torch.sigmoid(b.float()).unsqueeze(1) - - state_cula = torch.empty_like(state) - out_cula = torch.empty_like(v.squeeze(1)) - - def run_fla(): - return fused_recurrent_gated_delta_rule( - q=q, - k=k, - v=v, - g=gate, - beta=beta, - scale=config.head_k_dim**-0.5, - initial_state=state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - A_log=A_log, - dt_bias=dt_bias, - transpose_state_layout=False, - ) + def setup_cula(): + state_cula.copy_(x["state_kv"]) - def setup_cula() -> None: - state_cula.copy_(state) + def setup_sglang(): + state_sglang.copy_(x["state_vk"]) - def run_cula() -> None: + def run_cula(): cula_cuda.qwen35_layout_scalar_kda_decode( - mixed, - a, - b, - A_log, - dt_bias, - state_cula, - indices, - out_cula, + x["mixed_qkv"], x["a"], x["b"], x["A_log"], x["dt_bias"], + state_cula, x["indices"], out_cula, + ) + + def run_sglang(): + return sglang_kernel.packed_decode( + mixed_qkv=x["mixed_qkv"], a=x["a"], b=x["b"], + A_log=x["A_log"], dt_bias=x["dt_bias"], scale=k**-0.5, + ssm_states=state_sglang, cache_indices=x["indices"], + num_v_heads=hv, head_v_dim=v, ) - out_fla, state_fla = run_fla() setup_cula() run_cula() + setup_sglang() + out_sglang = run_sglang().squeeze(0) torch.cuda.synchronize() - out_rel_rms, out_max_abs = error_stats(out_fla.squeeze(1), out_cula) - state_rel_rms, state_max_abs = error_stats(state_fla, state_cula) + out_rrms = relative_rms(out_sglang, out_cula) + state_rrms = relative_rms(state_sglang, state_cula.transpose(-1, -2)) - fla_ms = benchmark_cuda(run_fla, warmup=args.warmup, rep=args.rep) + sglang_ms = benchmark_cuda(run_sglang, setup=setup_sglang, warmup=args.warmup, rep=args.rep) cula_ms = benchmark_cuda(run_cula, setup=setup_cula, warmup=args.warmup, rep=args.rep) return { "tokens": tokens, - "upstream_fla_ms": fla_ms, + "sglang_packed_ms": sglang_ms, "cula_fused_ms": cula_ms, - "speedup": fla_ms / cula_ms, - "out_rel_rms": out_rel_rms, - "out_max_abs": out_max_abs, - "state_rel_rms": state_rel_rms, - "state_max_abs": state_max_abs, + "speedup": sglang_ms / cula_ms, + "out_rel_rms": out_rrms, + "state_rel_rms": state_rrms, } def main() -> None: - parser = argparse.ArgumentParser(description="Upstream FLA vs fused cuLA Qwen3.5 decode") - parser.add_argument("--tokens", nargs="+", type=int, default=[1, 2, 4, 8, 16, 32, 64, 128]) + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--sglang-path", type=pathlib.Path, default=pathlib.Path("/sgl-workspace/sglang")) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--tokens", type=int, nargs="+", default=(1, 2, 4, 8, 16, 32, 64, 128)) parser.add_argument("--warmup", type=int, default=20) parser.add_argument("--rep", type=int, default=100) parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--tp-size", type=int, choices=[1, 2, 4, 8], default=1) parser.add_argument("--csv", type=pathlib.Path) args = parser.parse_args() if not torch.cuda.is_available(): raise RuntimeError("CUDA is required") - device = torch.device("cuda") + shape = load_shape(args.config_json, args.tp_size) + sglang_kernel = load_sglang(args.sglang_path) print( - f"device={torch.cuda.get_device_name(device)} torch={torch.__version__} " - f"cuda={torch.version.cuda} tp={args.tp_size} warmup/rep={args.warmup}/{args.rep}" + f"model={shape['model']} device={torch.cuda.get_device_name(0)} TP={args.tp_size} " + f"global_H/HV={shape['global_h']}/{shape['global_hv']} " + f"local_H/HV={shape['h']}/{shape['hv']} K/V={shape['k']}/{shape['v']}" ) - print("scope: upstream FLA recurrent operator vs cuLA fused Qwen layout + recurrent kernel") - print("| tokens | upstream_fla_ms | cula_fused_ms | speedup | out_rel_rms | state_rel_rms |") + print("state reset is outside timing; SGLang packed decode vs cuLA fused packed decode") + print("| tokens | sglang_packed_ms | cula_fused_ms | speedup | out_rrms | state_rrms |") print("|---:|---:|---:|---:|---:|---:|") - rows = [] for tokens in args.tokens: - row = run_case(tokens, args, device) + row = run_case(tokens, shape, sglang_kernel, args) rows.append(row) print( - f"| {tokens} | {row['upstream_fla_ms']:.4f} | {row['cula_fused_ms']:.4f} | " - f"{row['speedup']:.2f}x | {row['out_rel_rms']:.3e} | {row['state_rel_rms']:.3e} |" + f"| {tokens} | {row['sglang_packed_ms']:.4f} | {row['cula_fused_ms']:.4f} | " + f"{row['speedup']:.3f}x | {row['out_rel_rms']:.3e} | {row['state_rel_rms']:.3e} |" ) - - if args.csv is not None: + if args.csv: args.csv.parent.mkdir(parents=True, exist_ok=True) - with args.csv.open("w", newline="", encoding="utf-8") as file: - writer = csv.DictWriter(file, fieldnames=list(rows[0])) + with args.csv.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0])) writer.writeheader() writer.writerows(rows) - print(f"wrote {args.csv}") if __name__ == "__main__": From f7ee00afd9cd22e748d83c407883ca7a5e06ad32 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Sun, 2 Aug 2026 16:47:26 +0000 Subject: [PATCH 25/35] perf(qwen35): overlap long decode state load --- .../decode/qwen35_scalar_kda_kernel.hpp | 80 +++++++++++-------- 1 file changed, 47 insertions(+), 33 deletions(-) diff --git a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp index e84b4689..25e0aa7a 100644 --- a/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp +++ b/csrc/qwen35/decode/qwen35_scalar_kda_kernel.hpp @@ -719,7 +719,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( static_assert(kLocalQKHeads == Shape::kLocalQKHeads); static_assert(kHeadDimQK == 128); static_assert(kHeadDimV == 128); - static_assert(kTileV == 32 || kTileV == 64); + static_assert(kTileV == 32 || kTileV == 64 || kTileV == 128); static_assert(kHeadDimV % kTileV == 0); static_assert(kHeadDimQK % kThreads == 0); @@ -727,7 +727,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( __shared__ float k_smem[kHeadDimQK]; __shared__ float state_smem[kHeadDimQK][kTileV]; __shared__ float norm_smem[3]; - __shared__ float warp_reduce_smem[2 * kWarps]; + __shared__ float warp_reduce_smem[3 * kWarps]; __shared__ cutlass::arch::ClusterTransactionBarrier::ValueType state_barrier; const int hv_tile = static_cast(blockIdx.x); @@ -782,8 +782,46 @@ __global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( auto out_vec = gO(token_idx, hv, _); auto state_vk = gH_vk(state_row, hv, _, _); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Start the full-state transfer before the Q/K normalization and gate + // arithmetic. Those independent instructions hide a portion of the HBM + // latency for the long-token path. + if (tid == 0) { + cutlass::arch::ClusterTransactionBarrier::init(&state_barrier, 1); + cutlass::arch::ClusterTransactionBarrier::arrive_and_expect_tx( + &state_barrier, kHeadDimQK * kTileV * sizeof(float)); + } + __syncthreads(); + if constexpr (kTileV == kHeadDimV) { + constexpr int kStateBytes = kHeadDimQK * kHeadDimV * sizeof(float); + constexpr int kBulkChunkBytes = 32 * 1024; + constexpr int kBulkChunkFloats = kBulkChunkBytes / sizeof(float); + constexpr int kBulkChunks = kStateBytes / kBulkChunkBytes; + if (tid == 0) { +#pragma unroll + for (int chunk = 0; chunk < kBulkChunks; ++chunk) { + cp_async_bulk_shared_global( + &state_smem[0][0] + chunk * kBulkChunkFloats, + &state_vk(0, 0) + chunk * kBulkChunkFloats, + kBulkChunkBytes, + &state_barrier); + } + } + } else { +#pragma unroll 1 + for (int k_idx = tid; k_idx < kHeadDimQK; k_idx += kThreads) { + cp_async_bulk_shared_global( + &state_smem[k_idx][0], + &state_vk(v_tile * kTileV, k_idx), + kTileV * sizeof(float), + &state_barrier); + } + } +#endif + float q_norm_sq = 0.f; float k_norm_sq = 0.f; + float qk_raw_dot = 0.f; #pragma unroll for (int i = 0; i < kKPerThread; ++i) { const int k_idx = i * kThreads + tid; @@ -793,27 +831,32 @@ __global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( k_smem[k_idx] = k_raw; q_norm_sq += q_raw * q_raw; k_norm_sq += k_raw * k_raw; + qk_raw_dot += q_raw * k_raw; } q_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_norm_sq); k_norm_sq = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_norm_sq); + qk_raw_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_raw_dot); if (lane == 0) { warp_reduce_smem[warp_id] = q_norm_sq; warp_reduce_smem[kWarps + warp_id] = k_norm_sq; + warp_reduce_smem[2 * kWarps + warp_id] = qk_raw_dot; } __syncthreads(); if (warp_id == 0) { float q_block_sum = lane < kWarps ? warp_reduce_smem[lane] : 0.f; float k_block_sum = lane < kWarps ? warp_reduce_smem[kWarps + lane] : 0.f; + float qk_block_sum = lane < kWarps ? warp_reduce_smem[2 * kWarps + lane] : 0.f; q_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(q_block_sum); k_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(k_block_sum); + qk_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_block_sum); if (lane == 0) { norm_smem[0] = rsqrtf(q_block_sum + 1e-6f) * rsqrtf(static_cast(kHeadDimQK)); norm_smem[1] = rsqrtf(k_block_sum + 1e-6f); + norm_smem[2] = qk_block_sum * norm_smem[0] * norm_smem[1]; } } __syncthreads(); - float qk_dot = 0.f; #pragma unroll for (int i = 0; i < kKPerThread; ++i) { const int k_idx = i * kThreads + tid; @@ -821,21 +864,7 @@ __global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( const float k_normed = k_smem[k_idx] * norm_smem[1]; q_smem[k_idx] = q_normed; k_smem[k_idx] = k_normed; - qk_dot += q_normed * k_normed; - } - qk_dot = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_dot); - if (lane == 0) { - warp_reduce_smem[warp_id] = qk_dot; } - __syncthreads(); - if (warp_id == 0) { - float qk_block_sum = lane < kWarps ? warp_reduce_smem[lane] : 0.f; - qk_block_sum = Qwen35ScalarKdaDecodeMainloop::warp_sum(qk_block_sum); - if (lane == 0) { - norm_smem[2] = qk_block_sum; - } - } - __syncthreads(); const float a_val = static_cast(gA(token_idx, hv)); const float b_val = static_cast(gB(token_idx, hv)); @@ -845,20 +874,6 @@ __global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( const float beta = 1.f / (1.f + expf(-b_val)); #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - if (tid == 0) { - cutlass::arch::ClusterTransactionBarrier::init(&state_barrier, 1); - cutlass::arch::ClusterTransactionBarrier::arrive_and_expect_tx( - &state_barrier, kHeadDimQK * kTileV * sizeof(float)); - } - __syncthreads(); -#pragma unroll 1 - for (int k_idx = tid; k_idx < kHeadDimQK; k_idx += kThreads) { - cp_async_bulk_shared_global( - &state_smem[k_idx][0], - &state_vk(v_tile * kTileV, k_idx), - kTileV * sizeof(float), - &state_barrier); - } cutlass::arch::ClusterTransactionBarrier::wait(&state_barrier, 0); __syncthreads(); #else @@ -972,7 +987,6 @@ __global__ void qwen35_layout_scalar_kda_decode_long_vtile_kernel( state_vk(v_row, k_idx + 6) = state_new6; state_vk(v_row, k_idx + 7) = state_new7; } - } template @@ -990,7 +1004,7 @@ void launch_qwen35_layout_scalar_kda_decode_long_kernel( constexpr int kWarpTileV = 32; (void)kWarpTileV; if (token_count == 64 || token_count == 128) { - constexpr int kLongTileV = 64; + constexpr int kLongTileV = 128; dim3 grid(kLocalVHeads * (kHeadDimV / kLongTileV), token_count, 1); dim3 block(kLongTileV, 1, 1); qwen35_layout_scalar_kda_decode_long_vtile_kernel From 82b6cd61e144f6ad1c90de7e445b5a5c0a4412a9 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 15:31:43 +0000 Subject: [PATCH 26/35] perf(qwen35): optimize native-GVA scalar prefill on H200 --- csrc/kda/sm90/collective/load_tma.hpp | 18 +- csrc/kda/sm90/collective/mainloop_kda_fwd.hpp | 300 +++-- csrc/kda/sm90/collective/store_tma.hpp | 5 +- csrc/kda/sm90/kda_fwd_sm90.cu | 74 +- csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu | 47 + csrc/kda/sm90/kernel/builder_kda_fwd.hpp | 2 +- csrc/kda/sm90/kernel/kernel_kda_fwd.hpp | 36 +- csrc/kda/sm90/kernel/options.hpp | 3 + csrc/kda/sm90/kernel/tile_scheduler.hpp | 18 +- csrc/kda/sm90/prefill_kernel.hpp | 25 + csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh | 35 +- .../prefill/qwen35_scalar_kda_prefill.cu | 663 +++++++++- .../qwen35_scalar_kda_prefill_kernel.hpp | 1091 +++++++++++++++-- tests/test_qwen35_prefill.py | 182 +++ 14 files changed, 2279 insertions(+), 220 deletions(-) diff --git a/csrc/kda/sm90/collective/load_tma.hpp b/csrc/kda/sm90/collective/load_tma.hpp index 1d427b0c..f05dcbcd 100644 --- a/csrc/kda/sm90/collective/load_tma.hpp +++ b/csrc/kda/sm90/collective/load_tma.hpp @@ -106,22 +106,28 @@ struct CollectiveLoadTma { return g_full; } else if constexpr (kind == LoadKind::kAlpha) { // Alpha (gate) is per V/O head under GVA. + constexpr int AlphaWidth = decltype(size<1>(SmemLayout{}))::value; DPRINTF0_W( "slice view GMEM %s: seq_idx:%d head_idx:%d tok_offset:%lld\n", to_string(kind), work_desc.seq_idx, work_desc.o_head_idx(), work_desc.tok_offset); + constexpr bool ScalarAlpha = AlphaWidth == 4; + const int alpha_head_groups = + ScalarAlpha ? problem_size.num_v_heads / AlphaWidth : problem_size.num_v_heads; + const int alpha_head_group = + ScalarAlpha ? work_desc.o_head_idx() / AlphaWidth : work_desc.o_head_idx(); Tensor m_varlen_head = tma_load.get_tma_tensor(make_shape( problem_size.total_seqlen, - problem_size.head_size, - problem_size.num_v_heads)); // global view to the packed varlen sequence - Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx + Int{}, + alpha_head_groups)); // global view to packed tokens x 4-head groups + Tensor m_varlen = m_varlen_head(_, _, alpha_head_group); // slice group containing current V head Tensor m_offset = domain_offset( make_coord(work_desc.tok_offset, _0{}), m_varlen); // offset to start of the current sequence Tensor g_full = - local_tile(m_offset, make_tile(BlkSeqQ, HeadSize), make_coord(_, _0{})); // (blk, d, iter_blk) + local_tile(m_offset, make_tile(BlkSeqQ, Int{}), make_coord(_, _0{})); // (blk, d, iter_blk) return g_full; } else { // K lives in the QK head space; V lives in the V head space. @@ -140,8 +146,10 @@ struct CollectiveLoadTma { problem_size.total_seqlen, num_kv_heads)); // global view to the packed varlen sequence Tensor m_varlen = m_varlen_head(_, _, head_idx); // slice into current head_idx + const int feature_offset = + kIsK ? 0 : work_desc.value_tile_idx * HeadSize; Tensor m_offset = domain_offset( - make_coord(_0{}, work_desc.tok_offset), + make_coord(feature_offset, work_desc.tok_offset), m_varlen); // offset to start of the current sequence Tensor g_full = local_tile(m_offset, make_tile(HeadSize, BlkSeqKV), make_coord(_0{}, _)); // (d, blk, iter_blk) diff --git a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp index 66563e5a..4c7af669 100644 --- a/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp +++ b/csrc/kda/sm90/collective/mainloop_kda_fwd.hpp @@ -90,10 +90,12 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr bool kIsPersistent = find_option_t::value; static constexpr bool kInitStateFromInput = find_option_t::value; + static constexpr bool SplitValueDim = find_option_t::value; static constexpr int NumLoadWarpGroups = 1; - static constexpr int NumStateMmaWarpGroups = 2; + static constexpr int NumStateMmaWarpGroups = SplitValueDim ? 1 : 2; static constexpr int NumAuxMmaWarpGroups = 1; + static constexpr int NumValueTiles = SplitValueDim ? 2 : 1; static constexpr int StageCountQ = find_option_t, Options>::value; static constexpr int StageCountK = find_option_t, Options>::value; @@ -101,6 +103,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr int NeedsAlpha = find_option_t::value; static constexpr int NeedsBeta = find_option_t::value; + static constexpr bool ScalarAlpha = find_option_t::value; + static constexpr bool StateKVLayout = find_option_t::value; static_assert(NeedsAlpha && NeedsBeta, "Alpha and Beta are both used in KDA."); static constexpr int SafeGate = true; // only support safe_gate=true @@ -136,7 +140,11 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr auto BlkSeqKV = get<1>(TileShape{}); // Blk_K/V static constexpr auto HeadSize = get<2>(TileShape{}); // D (Dq, Dk, Dv all equal) static constexpr auto HeadSizeQK = HeadSize; - static constexpr auto HeadSizeV = HeadSize; + using HeadSizeVType = std::conditional_t< + SplitValueDim, + _64, + std::remove_cv_t>; + static constexpr auto HeadSizeV = HeadSizeVType{}; using HeadSizeHalf = _64; using HeadSizeQuar = _32; @@ -151,6 +159,11 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { using TileShapeO2 = decltype(make_shape(HeadSizeV, BlkSeqQ, BlkSeqKV)); using TileShapeO1 = decltype(make_shape(HeadSizeV, BlkSeqQ, HeadSizeQK)); + using StateMmaSchedule = std::conditional_t< + NumStateMmaWarpGroups == 1, + cutlass::gemm::KernelTmaWarpSpecialized, + cutlass::gemm::KernelTmaWarpSpecializedCooperative>; + static_assert(BlkSeqQ % 64 == 0); static_assert(BlkSeqQ == 64 || BlkSeqQ == 128); static_assert(BlkSeqQ == BlkSeqKV); @@ -200,24 +213,37 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeKV, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using SmemLayoutAlphaAtom = GMMA::Layout_K_SW128_Atom; - using SmemLayoutAlpha_SD = decltype(tile_to_shape( + using SmemLayoutAlphaVector_SD = decltype(tile_to_shape( SmemLayoutAlphaAtom{}, make_shape( shape<1>(TileShapeQK{}), shape<2>(TileShapeQK{}), Int{}))); // (blk_kv, head_size), (64, 128) - using GmemShapeAlpha = Shape; // (seqlen_k, d, h) + // TMA requires its first global mode to be contiguous. A single scalar + // head is strided by HV across tokens, so scalar mode loads four adjacent + // heads as one 16-byte unit and each CTA selects its hv % 4 lane. + using AlphaWidth = std::conditional_t(TileShapeQK{}))>; + using SmemLayoutAlphaScalar_SD = decltype(make_layout( + make_shape(shape<1>(TileShapeQK{}), _4{}, Int{}), + make_stride(_4{}, _1{}, Int<4 * size<1>(TileShapeQK{})>{}))); + using SmemLayoutAlphaLoad_SD = + std::conditional_t; + // Vector-alpha compute layouts stay intact for the generic path. Scalar + // specializations bypass them and fill MMA fragments from the compact + // [token, stage] tensor using token coordinates. + using SmemLayoutAlpha_SD = SmemLayoutAlphaVector_SD; + using GmemShapeAlpha = Shape; // (seqlen_k, d-or-1, h) using GmemStrideAlpha = Stride; using GmemLayoutAlpha = Layout; using GmemTiledCopyAlpha = cute::SM90_TMA_LOAD; using TMA_Alpha = decltype(make_tma_copy( GmemTiledCopyAlpha{}, make_tensor(make_gmem_ptr(static_cast(nullptr)), GmemLayoutAlpha{}), - take<0, 2>(SmemLayoutAlpha_SD{}), - select<1, 2>(TileShapeQK{}), + take<0, 2>(SmemLayoutAlphaLoad_SD{}), + make_shape(shape<1>(TileShapeQK{}), AlphaWidth{}), size<0>(ClusterShape{}))); // raw layout for copy @@ -247,7 +273,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeKV, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using RefLayoutKV = decltype(make_layout(select<0, 1>(TileShapeKV{}), LayoutRight{})); // (dv, dk) using CollectiveMmaO1 = typename cutlass::gemm::collective::CollectiveBuilder< @@ -263,7 +289,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeO1, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; // (blk_q,blk_k) to align with O2 mma, LayoutRight to align with QK mma output using DesiredLayoutQK = decltype(make_layout(select<0, 1>(TileShapeQK{}), LayoutRight{})); @@ -280,7 +306,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeO2, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using TiledMmaQK = typename CollectiveMmaQK::TiledMma; // Q@K^t using TiledMmaKV = decltype(convert_to_gmma_rs(typename CollectiveMmaKV::TiledMma{})); @@ -346,7 +372,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeSK, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; using ElementAccumulatorNewV = float; using TileShapeNewV = decltype(make_shape(HeadSizeV, BlkSeqKV, BlkSeqKV)); @@ -365,7 +391,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TileShapeNewV, ClusterShape, DummyStages, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + StateMmaSchedule>::CollectiveOp; // FIXME: K@K^t are not exactly the same as Q@K^t, but similar enough (what does this mean??) using TiledMmaKK = typename CollectiveMmaQK::TiledMma; // T = inv(I + strict_lower_triangular(K@K^t)) @@ -379,6 +405,9 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // only store the last row in Alpha using SmemLayoutAlphaLast = decltype(make_layout(make_shape(HeadSize, Int{}))); + using SmemLayoutAlphaLastScalar = decltype(make_layout(make_shape(_1{}, Int{}))); + using SmemLayoutAlphaLastStorage = + std::conditional_t; using SmemLayoutBeta = decltype(make_layout(make_shape(BlkSeqQ, Int{}))); using MainloopQPipeline = cutlass::PipelineTmaAsync; @@ -416,7 +445,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { static constexpr int LoadQBytes = size(QKSmemLayoutQ{}(_, _, _0{})) * sizeof(Element); static constexpr int LoadKBytes = size(KVSmemLayoutK{}(_, _, _0{})) * sizeof(Element); static constexpr int LoadVBytes = size(KVSmemLayoutV{}(_, _, _0{})) * sizeof(Element); - static constexpr int LoadAlphaBytes = size(QKQSmemLayoutAlpha{}(_, _, _0{})) * sizeof(ElementAlpha); + static constexpr int LoadAlphaBytes = size(SmemLayoutAlphaLoad_SD{}(_, _, _0{})) * sizeof(ElementAlpha); static constexpr int StoreOBytes = CollectiveStoreO::TmaTransactionBytes; using SharedStorageO = typename CollectiveStoreO::SharedStorage; @@ -429,7 +458,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { alignas( alignment_for_swizzle(KVSmemLayoutV{})) cute::array_aligned> smem_v; alignas(alignment_for_swizzle( - QKQSmemLayoutAlpha{})) cute::array_aligned> smem_alpha; + SmemLayoutAlphaLoad_SD{})) cute::array_aligned> smem_alpha; alignas( alignment_for_swizzle(SmemLayoutQK{})) cute::array_aligned> smem_qk; alignas(alignment_for_swizzle( @@ -442,7 +471,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { cute::array_aligned> smem_beta; // store last row in Alpha separately, used for S'=K^T NewV's epilogue and S+=decay(S') (one fused epilogue) - cute::array_aligned> smem_alpha_last; + cute::array_aligned> smem_alpha_last; }; using TMA_Q = typename CollectiveMmaQK::Params::TMA_A; @@ -454,7 +483,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { using LoadK = CollectiveLoadTma; using LoadV = CollectiveLoadTma; using LoadAlpha = - CollectiveLoadTma; + CollectiveLoadTma; using LoadBeta = CollectiveLoadVector< LoadKindVector::kBeta, @@ -493,6 +522,34 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { GmemLayoutBeta beta_layout; }; + template + CUTE_DEVICE static auto + make_state_tensor(Ptr ptr, ProblemShape const& problem_size, NumSeqs num_seqs) { + // The global state always remains a full [K=128,V=128] matrix. Split + // V64 CTAs select disjoint column tiles after constructing this view. + auto state_shape = make_shape( + Int{}, problem_size.head_size, problem_size.num_v_heads, num_seqs); + if constexpr (StateKVLayout) { + // Qwen exposes state as contiguous [N, HV, K, V]. Express that + // physical layout in the kernel's logical (K,V,HV,N) coordinates + // instead of paying for pre/post transpose kernels. + auto state_stride = make_stride( + int64_t(problem_size.head_size), + _1{}, + int64_t(HeadSizeQK) * problem_size.head_size, + int64_t(problem_size.num_v_heads) * int(HeadSizeQK) * problem_size.head_size); + return make_tensor(make_gmem_ptr(ptr), make_layout(state_shape, state_stride)); + } else { + return make_tensor(make_gmem_ptr(ptr), make_layout(state_shape, LayoutLeft{})); + } + } + + template + CUTE_DEVICE static auto + make_state_tensor(Ptr ptr, ProblemShape const& problem_size) { + return make_state_tensor(ptr, problem_size, problem_size.num_seqs); + } + template static bool can_implement(ProblemShape const& problem_size, Arguments const& args) { @@ -525,7 +582,9 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }, /*workspace=*/nullptr); - auto alpha_shape = make_shape(s, d, problem_size.num_v_heads); + const int32_t alpha_head_groups = + ScalarAlpha ? problem_size.num_v_heads / int(AlphaWidth{}) : problem_size.num_v_heads; + auto alpha_shape = make_shape(s, AlphaWidth{}, alpha_head_groups); auto alpha_stride = make_stride( get<0>(args.dAlpha), // seqlen stride get<1>(args.dAlpha), // head_dim stride @@ -535,8 +594,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { TMA_Alpha tma_load_alpha = make_tma_copy( GmemTiledCopyAlpha{}, mAlpha, - take<0, 2>(SmemLayoutAlpha_SD{}), - select<1, 2>(TileShapeQK{}), + take<0, 2>(SmemLayoutAlphaLoad_SD{}), + make_shape(shape<1>(TileShapeQK{}), AlphaWidth{}), size<0>(ClusterShape{})); auto params_kv_v = CollectiveMmaKV_G2S::to_underlying_arguments( @@ -617,10 +676,11 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto k_collective_load = LoadK(params.tma_load_k, k_pipeline, storage.smem_k); auto v_collective_load = LoadV(params.tma_load_v, v_pipeline, storage.smem_v); auto alpha_collective_load = LoadAlpha{params.tma_load_alpha, alpha_pipeline, storage.smem_alpha}; + auto v_load_tile_shape = make_shape(BlkSeqQ, BlkSeqKV, HeadSizeV); auto q_src_dst = q_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); auto k_src_dst = k_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); - auto v_src_dst = v_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); + auto v_src_dst = v_collective_load.partition_SD(problem_size, v_load_tile_shape, work_desc); auto alpha_src_dst = alpha_collective_load.partition_SD(problem_size, load_tile_shape, work_desc); CUTE_NO_UNROLL @@ -673,7 +733,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarp; Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); - Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); + Tensor sAlphaLoad = make_tensor(make_smem_ptr(storage.smem_alpha.data()), SmemLayoutAlphaLoad_SD{}); + Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLastStorage{}); auto extract_loop_body = [&](int blk, auto is_final_block_) INLINE_LAMBDA { constexpr bool is_final_block = decltype(is_final_block_)::value; @@ -681,15 +742,22 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { int B = is_final_block ? valid_seq_len(work_desc, blk) : BlkSeqKV; auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); + auto sAlphaLoadCurr = sAlphaLoad(_, _, alpha_smem_pipe_read.index()); Tensor sAlast_out = sAlast(_, alpha_last_smem_pipe_write.index()); alpha_pipeline.consumer_wait(alpha_smem_pipe_read); alpha_last_pipeline.producer_acquire(alpha_last_smem_pipe_write); // each thread copy 4 elements, total 128 elements with one warp - CUTE_UNROLL - for (int t = thread_idx; t < HeadSize; t += 32) { - sAlast_out(t) = sAqkq_curr(B - 1, t); + if constexpr (ScalarAlpha) { + if (thread_idx == 0) { + sAlast_out(_0{}) = sAlphaLoadCurr(B - 1, work_desc.o_head_idx() & 3); + } + } else { + CUTE_UNROLL + for (int t = thread_idx; t < HeadSize; t += 32) { + sAlast_out(t) = sAqkq_curr(B - 1, t); + } } cutlass::arch::fence_view_async_shared(); @@ -775,10 +843,13 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { Tensor Beta = make_tensor(make_smem_ptr(storage.smem_beta.data()), SmemLayoutBeta{}); Tensor AlphaLast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); + Tensor AlphaLastStorage = + make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLastStorage{}); Tensor sQqk = make_tensor(make_smem_ptr(storage.smem_q.data()), QKSmemLayoutQ{}); Tensor sKqk = make_tensor(make_smem_ptr(storage.smem_k.data()), QKSmemLayoutK{}); Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); + Tensor sAlphaLoad = make_tensor(make_smem_ptr(storage.smem_alpha.data()), SmemLayoutAlphaLoad_SD{}); Tensor sVkv = make_tensor(make_smem_ptr(storage.smem_v.data()), KVSmemLayoutV{}); Tensor sQK = make_tensor(make_smem_ptr(storage.smem_qk.data()), SmemLayoutQK{}); Tensor sO = make_tensor(make_smem_ptr(storage.smem_o.data()), SmemLayoutO{}); @@ -883,7 +954,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto thr_copy_o = tiled_copy_o.get_thread_slice(thread_idx); auto tOsO = thr_copy_o.partition_D(sO); - auto const cO = make_identity_tensor(Shape, Int>{}); + auto const cO = make_identity_tensor(Shape, Int>{}); Tensor tOcO = o1_thr_mma.partition_C(cO); auto const seq_idx = work_desc.seq_idx; @@ -902,12 +973,13 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto kv_load = [&](auto& tKVrKV) INLINE_LAMBDA { DPRINTF0_WG("[%d,%d,%d,%d]>> load tKVgKV -> tKVrKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); // GVA: state is stored per V/O head. - int num_state_heads = problem_size.num_v_heads; int state_head_idx = work_desc.o_head_idx(); - auto gKV = make_tensor( - make_gmem_ptr(params.ptr_input_state), - make_layout(make_shape(Int{}, Int{}, num_state_heads, problem_size.num_seqs)))( - _, _, state_head_idx, seq_idx); // (KDim, VDim), K-contiguous + auto gKV_full = make_state_tensor(params.ptr_input_state, problem_size)( + _, _, state_head_idx, seq_idx); // full (KDim, VDim) + auto gKV = local_tile( + gKV_full, + make_tile(Int{}, HeadSizeVType{}), + make_coord(_0{}, work_desc.value_tile_idx)); auto tiled_copy_kv = make_tiled_copy_C(Copy_Atom{}, kv_tiled_mma); auto thr_copy_kv = tiled_copy_kv.get_thread_slice(thread_idx); @@ -944,12 +1016,13 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { } DPRINTF0_WG("[%d,%d,%d,%d]>> save tKVrKV -> tKVgKV\n", seq_idx, q_head_idx, k_head_idx, v_head_idx); // GVA: state is stored per V/O head. - int num_state_heads = problem_size.num_v_heads; int state_head_idx = work_desc.o_head_idx(); - auto gKV = make_tensor( - make_gmem_ptr(params.ptr_output_state), - make_layout(make_shape(Int{}, Int{}, num_state_heads, out_num_seqs)))( - _, _, state_head_idx, out_seq_idx); // (KDim, VDim), K-contiguous + auto gKV_full = make_state_tensor(params.ptr_output_state, problem_size, out_num_seqs)( + _, _, state_head_idx, out_seq_idx); // full (KDim, VDim) + auto gKV = local_tile( + gKV_full, + make_tile(Int{}, HeadSizeVType{}), + make_coord(_0{}, work_desc.value_tile_idx)); auto tiled_copy_kv = make_tiled_copy_C(Copy_Atom{}, kv_tiled_mma); auto thr_copy_kv = tiled_copy_kv.get_thread_slice(thread_idx); @@ -959,12 +1032,17 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }; auto s_decay = [&](auto& tKVrKV, auto const& alpha_last_smem_pipe_read) INLINE_LAMBDA { - Tensor alpha_last_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); - for_each(make_int_sequence{}, [&](auto i) { - auto coord = tKVcS(i); - auto [s, t] = coord; // (head_size_v, head_size_k) - tKVrKV(i) *= exp2f(alpha_last_curr(t)); - }); + if constexpr (ScalarAlpha) { + const float decay = exp2f(AlphaLastStorage(_0{}, alpha_last_smem_pipe_read.index())); + for_each(make_int_sequence{}, [&](auto i) { tKVrKV(i) *= decay; }); + } else { + Tensor alpha_last_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); + for_each(make_int_sequence{}, [&](auto i) { + auto coord = tKVcS(i); + auto [s, t] = coord; // (head_size_v, head_size_k) + tKVrKV(i) *= exp2f(alpha_last_curr(t)); + }); + } }; auto o1_epi = [&](auto& tOrO1) INLINE_LAMBDA { @@ -1034,6 +1112,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto sK_scaled_curr = sQ_K_scaled(_, _, _1{}); auto sAlast_curr = AlphaLast(_, alpha_last_smem_pipe_read.index()); auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); + auto sAlphaLoadCurr = sAlphaLoad(_, _, alpha_smem_pipe_read.index()); auto sQqk_slice = flat_divide(sQqk_curr, tiler_qk); auto sKqk_slice = flat_divide(sKqk_curr, tiler_qk); auto sQ_scaled_slice = flat_divide(sQ_scaled_curr, tiler_qk); @@ -1054,12 +1133,12 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { if constexpr (!is_first_block) { // make sure sQ_K_scaled is already consumed for previous K^@V cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // Each WG iterates over 2 slices of 32 elements each. - // WG0 (thread_idx < 128): wg_idx=0, processes alpha indices {0,1}, Q/K dim1=0 - // WG1 (thread_idx >= 128): wg_idx=1, processes alpha indices {2,3}, Q/K dim1=1 + // Divide the four 32-d Q/K quarters over the active state WGs. + // V128 uses two quarters per WG; split V64 uses one WG for all four. { - int wg_idx = thread_idx / 128; // 0 or 1 - int alpha_base = wg_idx * 2; // 0 or 2 + constexpr int kQuarterCount = int(HeadSizeQK) / int(HeadSizeQuar{}); + constexpr int kQuartersPerWG = kQuarterCount / NumStateMmaWarpGroups; + int wg_idx = thread_idx / 128; // Allocate Q/K register fragments once (reused across slices) // Only shape/layout matters for partition_fragment_A, use compile-time indices @@ -1069,17 +1148,26 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { qk_thr_mma_rs_quar.partition_fragment_A(sKqk_slice(_, _, _0{}, make_coord(_0{}, _0{}))); auto tArA = make_fragment_like(tQKrQ_wg); - for (int s = 0; s < 2; ++s) { - // S2R Alpha: alpha_col = wg_idx * 2 + s - int alpha_col = alpha_base + s; - auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); - auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); - copy(CopyAlphaAtom{}, tAsA_cur, tArA); + for (int local_quarter = 0; local_quarter < kQuartersPerWG; ++local_quarter) { + int quarter = wg_idx * kQuartersPerWG + local_quarter; + int quarter_col = quarter & 1; + int quarter_group = quarter >> 1; + int alpha_col = quarter; + if constexpr (ScalarAlpha) { + for_each(make_int_sequence{}, [&](auto i) { + auto [seq, _] = tQcMq_quar(i); + tArA(i) = sAlphaLoadCurr(seq, work_desc.o_head_idx() & 3); + }); + } else { + auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); + auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); + copy(CopyAlphaAtom{}, tAsA_cur, tArA); + } cute::transform(tArA, [](auto g) { return exp2f(g); }); // S2R Q - auto sQqk_cur = sQqk_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sQqk_cur = sQqk_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsQ_cur = thr_load_qk_quar.partition_S(sQqk_cur); auto tQKrQ_cv = thr_load_qk_quar.retile_D(tQKrQ_wg); copy(tiled_load_qk_quar, tQKsQ_cur, tQKrQ_cv); @@ -1091,13 +1179,14 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }); // R2S Q -> stage 0 - auto sQ_scaled_cur = sQ_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sQ_scaled_cur = + sQ_scaled_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsQ_out = thr_store_qk_quar.partition_D(sQ_scaled_cur); auto tQKrQ_out_cv = thr_store_qk_quar.retile_S(tQKrQ_wg); copy(tiled_store_qk_quar, tQKrQ_out_cv, tQKsQ_out); // S2R K - auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_cur = thr_load_qk_quar.partition_S(sKqk_cur); auto tQKrK_cv = thr_load_qk_quar.retile_D(tQKrK_wg); copy(tiled_load_qk_quar, tQKsK_cur, tQKrK_cv); @@ -1109,7 +1198,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }); // R2S K -> stage 1 - auto sK_scaled_cur = sK_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sK_scaled_cur = + sK_scaled_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_out = thr_store_qk_quar.partition_D(sK_scaled_cur); auto tQKrK_out_cv = thr_store_qk_quar.retile_S(tQKrK_wg); copy(tiled_store_qk_quar, tQKrK_out_cv, tQKsK_out); @@ -1286,40 +1376,58 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // synchronize 2 WGs before rewriting sQ_K_scaled cutlass::arch::NamedBarrier::arrive_and_wait(NumStateMmaThreads, KdaNamedBarriers::StateMath); - // exp(alpha_last - alpha) * K - // Each WG iterates over 2 slices of 32 elements each. - // WG0 (thread_idx < 128): wg_idx=0, alpha_last indices {0,1}, K/output dim1=0 - // WG1 (thread_idx >= 128): wg_idx=1, alpha_last indices {2,3}, K/output dim1=1 + // exp(alpha_last - alpha) * K over all four 32-d quarters. { - int wg_idx = thread_idx / 128; // 0 or 1 - int alpha_base = wg_idx * 2; // 0 or 2 + constexpr int kQuarterCount = int(HeadSizeQK) / int(HeadSizeQuar{}); + constexpr int kQuartersPerWG = kQuarterCount / NumStateMmaWarpGroups; + int wg_idx = thread_idx / 128; // Allocate K/Alpha register fragments once (reused across slices) auto tQKrK_wg = qk_thr_mma_rs_quar.partition_fragment_A(sKqk_slice(_, _, _0{}, make_coord(_0{}, _0{}))); auto tArA_wg = make_fragment_like(tQKrK_wg); + float scalar_alpha_last = 0.0f; + if constexpr (ScalarAlpha) { + scalar_alpha_last = AlphaLastStorage(_0{}, alpha_last_smem_pipe_read.index()); + } - for (int s = 0; s < 2; ++s) { + for (int local_quarter = 0; local_quarter < kQuartersPerWG; ++local_quarter) { + int quarter = wg_idx * kQuartersPerWG + local_quarter; + int quarter_col = quarter & 1; + int quarter_group = quarter >> 1; // S2R Alpha - int alpha_col = alpha_base + s; - auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); - auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); - copy(CopyAlphaAtom{}, tAsA_cur, tArA_wg); + int alpha_col = quarter; + if constexpr (ScalarAlpha) { + for_each(make_int_sequence{}, [&](auto i) { + auto [seq, _] = tQcMq_quar(i); + tArA_wg(i) = sAlphaLoadCurr(seq, work_desc.o_head_idx() & 3); + }); + } else { + auto sA_cur = sAqkq_slice(_, _, _0{}, make_coord(0, alpha_col)); + auto tAsA_cur = qk_thr_mma_rs_quar.partition_A(sA_cur); + copy(CopyAlphaAtom{}, tAsA_cur, tArA_wg); + } // S2R K - auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sKqk_cur = sKqk_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_cur = thr_load_qk_quar.partition_S(sKqk_cur); auto tQKrK_cv = thr_load_qk_quar.retile_D(tQKrK_wg); copy(tiled_load_qk_quar, tQKsK_cur, tQKrK_cv); // element-wise: exp(alpha_last - alpha) * K - int alast_idx = alpha_base + s; + int alast_idx = quarter; auto alpha_last_cur = sAlast_slice(_, alast_idx); for_each(make_int_sequence{}, [&](auto i) { auto coord = tQcMq_quar(i); auto [seq, t] = coord; auto alpha = tArA_wg(i); auto k = tQKrK_wg(i); - auto alpha_last = alpha_last_cur(t); + auto alpha_last = [&]() { + if constexpr (ScalarAlpha) { + return scalar_alpha_last; + } else { + return alpha_last_cur(t); + } + }(); auto k_scaled = Element(exp2f(alpha_last - alpha) * float(k)); tQKrK_wg(i) = k_scaled; if constexpr (is_final_block) { @@ -1330,7 +1438,8 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { }); // R2S K -> stage 0 (reuse for KV update) - auto sQ_scaled_cur = sQ_scaled_slice(_, _, _0{}, make_coord(s, wg_idx)); + auto sQ_scaled_cur = + sQ_scaled_slice(_, _, _0{}, make_coord(quarter_col, quarter_group)); auto tQKsK_out = thr_store_qk_quar.partition_D(sQ_scaled_cur); auto tQKrK_out_cv = thr_store_qk_quar.retile_S(tQKrK_wg); copy(tiled_store_qk_quar, tQKrK_out_cv, tQKsK_out); @@ -1444,6 +1553,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { Tensor sAqkq = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKQSmemLayoutAlpha{}); Tensor sAqkk = make_tensor(make_smem_ptr(storage.smem_alpha.data()), QKKSmemLayoutAlpha{}); + Tensor sAlphaLoad = make_tensor(make_smem_ptr(storage.smem_alpha.data()), SmemLayoutAlphaLoad_SD{}); Tensor sAlast = make_tensor(make_smem_ptr(storage.smem_alpha_last.data()), SmemLayoutAlphaLast{}); Tensor sKkv = make_tensor(make_smem_ptr(storage.smem_k.data()), KVSmemLayoutK{}); @@ -1517,6 +1627,10 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // index tensor auto cMqk_subchunk = make_identity_tensor(select<0, 1>(TileShape_SubChunk{})); auto tQKcMqk_subchunk = thr_mma_subchunk.partition_C(cMqk_subchunk); + auto cMq_subchunk = make_identity_tensor(select<0, 2>(TileShape_SubChunk{})); + auto tQcMq_bf16_subchunk = thr_mma_bf16_subchunk.partition_A(cMq_subchunk); + auto cNk_subchunk = make_identity_tensor(select<1, 2>(TileShape_SubChunk{})); + auto tKcNk_bf16_subchunk = thr_mma_bf16_subchunk.partition_B(cNk_subchunk); // do MMA at the granularity of 16x16x64 with two warps constexpr auto tiler_subchunk_alpha = Shape<_16, Shape<_32, _1>>{}; @@ -1525,6 +1639,7 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto sQqk_curr = sQqk(_, _, q_smem_pipe_read.index()); auto sKqk_curr = sKqk(_, _, k_smem_pipe_read.index()); auto sAqkq_curr = sAqkq(_, _, alpha_smem_pipe_read.index()); + auto sAlphaLoadCurr = sAlphaLoad(_, _, alpha_smem_pipe_read.index()); Tensor sBeta_curr = Beta(_, beta_smem_pipe_read.index()); // (_16,(_32,_1),_4,(_2,_2)):(_64,(_1,_0),_1024,(_32,_4096)) @@ -1565,11 +1680,20 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // layout) auto s2r_compute_subchunk_operandA = [&](auto r_, int j, int j0, int j1) INLINE_LAMBDA { // S2R g_r_j in BF16 MMA operand A layout (single load) - Tensor sAqkq_r_j = sAqkq_slice(_, _, r_, make_coord(_0{}, j)); - Tensor tAsA_r_j = alpha_Q_bf16_thr_copy.partition_S(sAqkq_r_j); Tensor tArA_r_j = make_fragment_like(tv_layout_bf16_mma_A); - Tensor tArA_r_j_cv = alpha_Q_bf16_thr_copy.retile_D(tArA_r_j); - copy(alpha_Q_bf16_tiled_copy, tAsA_r_j, tArA_r_j_cv); + if constexpr (ScalarAlpha) { + constexpr int kSubchunkRows = size<0>(TileShape_SubChunk{}); + for_each(make_int_sequence{}, [&](auto i) { + auto [row, _] = tQcMq_bf16_subchunk(i); + tArA_r_j(i) = + sAlphaLoadCurr(int(r_) * kSubchunkRows + int(row), work_desc.o_head_idx() & 3); + }); + } else { + Tensor sAqkq_r_j = sAqkq_slice(_, _, r_, make_coord(_0{}, j)); + Tensor tAsA_r_j = alpha_Q_bf16_thr_copy.partition_S(sAqkq_r_j); + Tensor tArA_r_j_cv = alpha_Q_bf16_thr_copy.retile_D(tArA_r_j); + copy(alpha_Q_bf16_tiled_copy, tAsA_r_j, tArA_r_j_cv); + } // Derive g_first (alpha[row=0, :]) from tArA_r_j via warp shuffle, // directly into operand B layout (8 values instead of 16). @@ -1577,7 +1701,14 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { // v1=0 subset of operand A. We shuffle v1=0 values from t1=0 thread and // output directly as operand B fragment, saving 8 float registers. Tensor tArAfirst_r_j_kt = make_fragment_like(tv_layout_bf16_mma_B); - broadcast_row0_operandA_to_operandB_bf16_layout(tArA_r_j, tArAfirst_r_j_kt, local_thread_idx); + if constexpr (ScalarAlpha) { + const float g_first = sAlphaLoadCurr( + int(r_) * size<0>(TileShape_SubChunk{}), work_desc.o_head_idx() & 3); + fill(tArAfirst_r_j_kt, g_first); + } else { + broadcast_row0_operandA_to_operandB_bf16_layout( + tArA_r_j, tArAfirst_r_j_kt, local_thread_idx); + } // gqn_r_j = exp2(g_r_j - g_r_j_first[None, :]) in BF16 MMA A layout. // g_first per k-iter is in tArAfirst_r_j_kt: frag_B(2j)=K_lo, frag_B(2j+1)=K_hi. @@ -1635,11 +1766,20 @@ struct FlatMainloopTmaWarpSpecializedKdaFwd { auto s2r_compute_subchunk_operandB = [&](auto c_, int j, int j0, int j1, auto const& tArAfirst_kt) INLINE_LAMBDA { // S2R g_c_j in BF16 MMA operand B layout - Tensor sAqkq_c_j = sAqkq_slice(_, _, c_, make_coord(_0{}, j)); - Tensor tAsA_c_j = alpha_Kt_bf16_thr_copy.partition_S(sAqkq_c_j); Tensor tArA_c_j = make_fragment_like(tv_layout_bf16_mma_B); - Tensor tArA_c_j_cv = alpha_Kt_bf16_thr_copy.retile_D(tArA_c_j); - copy(alpha_Kt_bf16_tiled_copy, tAsA_c_j, tArA_c_j_cv); + if constexpr (ScalarAlpha) { + constexpr int kSubchunkCols = size<1>(TileShape_SubChunk{}); + for_each(make_int_sequence{}, [&](auto i) { + auto [col, _] = tKcNk_bf16_subchunk(i); + tArA_c_j(i) = + sAlphaLoadCurr(int(c_) * kSubchunkCols + int(col), work_desc.o_head_idx() & 3); + }); + } else { + Tensor sAqkq_c_j = sAqkq_slice(_, _, c_, make_coord(_0{}, j)); + Tensor tAsA_c_j = alpha_Kt_bf16_thr_copy.partition_S(sAqkq_c_j); + Tensor tArA_c_j_cv = alpha_Kt_bf16_thr_copy.retile_D(tArA_c_j); + copy(alpha_Kt_bf16_tiled_copy, tAsA_c_j, tArA_c_j_cv); + } // compute gktn_c_j = exp2(g_first - g_c_j) in BF16 MMA B layout cute::transform( diff --git a/csrc/kda/sm90/collective/store_tma.hpp b/csrc/kda/sm90/collective/store_tma.hpp index 0f7f7c1a..e2349f6a 100644 --- a/csrc/kda/sm90/collective/store_tma.hpp +++ b/csrc/kda/sm90/collective/store_tma.hpp @@ -184,7 +184,6 @@ struct CollectiveStoreTma { CUTE_DEVICE auto partition_SD(ProblemSize const& problem_size, TileShape const& tile_shape, WorkDesc const& work_desc) { constexpr auto BlkSeqQ = decltype(get<0>(tile_shape))::value; - constexpr auto HeadSize = decltype(get<2>(tile_shape))::value; Tensor g = [&] { DPRINTF0_W( @@ -198,10 +197,10 @@ struct CollectiveStoreTma { problem_size.num_v_heads)); // O lives in the V/O head space under GVA Tensor m_varlen = m_varlen_head(_, _, work_desc.o_head_idx()); // slice into current head_idx Tensor m_offset = domain_offset( - make_coord(_0{}, work_desc.tok_offset), + make_coord(work_desc.value_tile_idx * int(SizeM{}), work_desc.tok_offset), m_varlen); // offset to start of the current sequence Tensor g_full = - local_tile(m_offset, make_tile(HeadSize, BlkSeqQ), make_coord(_0{}, _)); // (d, blk, iter_blk) + local_tile(m_offset, make_tile(SizeM{}, BlkSeqQ), make_coord(_0{}, _)); // (d, blk, iter_blk) return g_full; }(); Tensor s = make_tensor(make_smem_ptr(storage_.data()), SmemLayoutO{}); diff --git a/csrc/kda/sm90/kda_fwd_sm90.cu b/csrc/kda/sm90/kda_fwd_sm90.cu index ed855db6..48dfaaac 100644 --- a/csrc/kda/sm90/kda_fwd_sm90.cu +++ b/csrc/kda/sm90/kda_fwd_sm90.cu @@ -22,6 +22,7 @@ namespace kda::sm90 { using namespace cute; +using bf16 = cute::bfloat16_t; // Forward declaration of the per-variant launcher (defined in .cuh, instantiated in separate TUs) template < @@ -33,7 +34,10 @@ template < typename TO, typename TQKV, typename TState, - typename TBeta = float> + typename TBeta = float, + bool ScalarAlpha = false, + bool StateKVLayout = false, + bool SplitValueDim = false> void launch_kda_fwd_prefill_kernel_gbai( cudaStream_t stream, @@ -58,6 +62,72 @@ launch_kda_fwd_prefill_kernel_gbai( int32_t const* raw_cu_seqlens, int32_t raw_num_seqs); +void +launch_qwen35_scalar_kda_fwd_prefill_kernel( + cudaStream_t stream, + void* output, + float* output_state, + void const* q, + void const* k, + void const* v, + float const* input_state, + float const* alpha, + float const* beta, + int32_t const* cu_seqlens, + uint8_t* workspace_buffer, + int32_t num_seqs, + int32_t num_qk_heads, + int32_t num_v_heads, + int32_t head_size, + int64_t total_seqlen, + float scale, + bool has_initial_state, + int32_t sm_count) { + if (has_initial_state) { + launch_kda_fwd_prefill_kernel_gbai< + true, true, true, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + stream, + static_cast(output), + output_state, + static_cast(q), + static_cast(k), + static_cast(v), + input_state, + alpha, + beta, + cu_seqlens, + workspace_buffer, + num_seqs, + num_qk_heads, + num_v_heads, + head_size, + total_seqlen, + scale, + sm_count); + } else { + launch_kda_fwd_prefill_kernel_gbai< + true, true, false, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + stream, + static_cast(output), + output_state, + static_cast(q), + static_cast(k), + static_cast(v), + nullptr, + alpha, + beta, + cu_seqlens, + workspace_buffer, + num_seqs, + num_qk_heads, + num_v_heads, + head_size, + total_seqlen, + scale, + sm_count); + } +} + template < typename ArchTag, // TODO: hide this typename TO, @@ -132,8 +202,6 @@ launch_kda_fwd_prefill_kernel( #undef LAUNCH } -using bf16 = cute::bfloat16_t; - // TBeta=float (default) template void launch_kda_fwd_prefill_kernel( diff --git a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu index 0da2986e..f0a51a92 100644 --- a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu +++ b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu @@ -63,4 +63,51 @@ INSTANTIATE_GBAI(true, true, true, true, bf16); #undef INSTANTIATE_GBAI +// Qwen scalar-G specialization: compact alpha [T, HV] and external state +// buffers in contiguous [K, V] layout. Keep this separate from the generic +// vector-alpha instantiations above. +template void +launch_kda_fwd_prefill_kernel_gbai< + true, true, false, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + cudaStream_t, + bf16*, + float*, + bf16 const*, + bf16 const*, + bf16 const*, + float const*, + float const*, + float const*, + int32_t const*, + uint8_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int64_t, + float, + int32_t); + +template void +launch_kda_fwd_prefill_kernel_gbai< + true, true, true, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( + cudaStream_t, + bf16*, + float*, + bf16 const*, + bf16 const*, + bf16 const*, + float const*, + float const*, + float const*, + int32_t const*, + uint8_t*, + int32_t, + int32_t, + int32_t, + int32_t, + int64_t, + float, + int32_t); + } // namespace kda::sm90 diff --git a/csrc/kda/sm90/kernel/builder_kda_fwd.hpp b/csrc/kda/sm90/kernel/builder_kda_fwd.hpp index 74e9c43a..cc9a606f 100644 --- a/csrc/kda/sm90/kernel/builder_kda_fwd.hpp +++ b/csrc/kda/sm90/kernel/builder_kda_fwd.hpp @@ -70,7 +70,7 @@ struct FlatBuilderKdaFwd< static constexpr bool kIsPersistent = find_option_t::value; static_assert(!kIsPersistent, "not implemented"); - using TileScheduler = kda::sm90::kernel::IndividualTileScheduler; + using TileScheduler = kda::sm90::kernel::IndividualTileScheduler; // using TileScheduler = std::conditional_t; diff --git a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp b/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp index ac597f7b..ef326859 100644 --- a/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp +++ b/csrc/kda/sm90/kernel/kernel_kda_fwd.hpp @@ -198,7 +198,11 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { get_register_requirements(MaxThreadsPerBlock, MinBlocksPerMultiprocessor, NumStateMmaWarpGroups); static constexpr uint32_t LdStRegisterRequirement = get<0>(RegisterRequirements); static constexpr uint32_t StateMmaRegisterRequirement = get<1>(RegisterRequirements); - static constexpr uint32_t AuxMmaRegisterRequirement = get<2>(RegisterRequirements); + // The V64 specialization statically uses 168 registers/thread. setmaxnreg + // `.inc 152` is illegal when the requested aux ceiling is below that + // initial allocation, so keep the aux WG slightly above the static value. + static constexpr uint32_t AuxMmaRegisterRequirement = + NumStateMmaWarpGroups == 1 ? 176 : get<2>(RegisterRequirements); static size_t get_workspace_size(Arguments const& args) { @@ -236,13 +240,6 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { CUTE_DEVICE void operator()(const Params& params, char* smem) { - enum class WarpGroupRole { - LdSt = 0, - Math0 = 1, - Math1 = 2, - MathA = 3, // auxiliary math WG - }; - // NOTE: CollectiveInverse will have more utilization on warp 0&1 // so we put beta and alpha preprocessing on warp 2&3 enum class LdStWarpRole { @@ -261,7 +258,10 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { int warp_idx = cutlass::canonical_warp_idx_sync(); int warp_idx_in_wg = warp_idx % cutlass::NumWarpsPerWarpGroup; int warp_group_idx = cutlass::canonical_warp_group_idx(); - auto warp_group_role = WarpGroupRole(warp_group_idx); + bool is_load_wg = warp_group_idx == 0; + bool is_state_wg = + warp_group_idx >= 1 && warp_group_idx < 1 + NumStateMmaWarpGroups; + bool is_aux_wg = warp_group_idx == 1 + NumStateMmaWarpGroups; auto ldst_warp_role = LdStWarpRole(warp_idx_in_wg); int lane_predicate = cute::elect_one_sync(); @@ -323,7 +323,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { OrderedMathBarriers math_barriers; - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadQKV) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::LoadQKV) { DPRINTF0_W("ldst_warp_role: LoadQKV Alpha\n"); q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Producer; k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Producer; @@ -332,23 +332,23 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Producer; } } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::StoreO) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::StoreO) { DPRINTF0_W("ldst_warp_role: StoreO\n"); o_pipeline_params.role = MainloopOPipeline::ThreadCategory::Consumer; } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadBeta) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::LoadBeta) { if constexpr (NeedsBeta) { beta_pipeline_params.role = MainloopBetaPipeline::ThreadCategory::Producer; } } - if (warp_group_role == WarpGroupRole::LdSt && ldst_warp_role == LdStWarpRole::LoadAlpha) { + if (is_load_wg && ldst_warp_role == LdStWarpRole::LoadAlpha) { // LoadAlpha warp consumes alpha_pipeline (reads last row) and produces alpha_last_pipeline if constexpr (NeedsAlpha) { alpha_pipeline_params.role = MainloopAlphaPipeline::ThreadCategory::Consumer; } alpha_last_pipeline_params.role = MainloopAlphaLastPipeline::ThreadCategory::Producer; } - if (warp_group_role == WarpGroupRole::Math0 || warp_group_role == WarpGroupRole::Math1) { + if (is_state_wg) { DPRINTF0_WG("warp_group_role: MathX\n"); q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Consumer; k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Consumer; @@ -368,7 +368,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { math_barriers.init(warp_group_idx - 1); } - if (warp_group_role == WarpGroupRole::MathA) { + if (is_aux_wg) { DPRINTF0_WG("warp_group_role: MathA\n"); q_pipeline_params.role = MainloopQPipeline::ThreadCategory::Consumer; k_pipeline_params.role = MainloopKPipeline::ThreadCategory::Consumer; @@ -453,7 +453,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { CollectiveMainloop collective_mainloop; - if (warp_group_role == WarpGroupRole::LdSt) { + if (is_load_wg) { DPRINTF0_WG("LsSt warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, LdStRegisterRequirement); cutlass::arch::warpgroup_reg_dealloc(); if (ldst_warp_role == LdStWarpRole::LoadQKV) { @@ -549,7 +549,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { o_smem_pipe_read, storage.tensors.mainloop.smem_o); } - } else if (warp_group_role == WarpGroupRole::Math0 || warp_group_role == WarpGroupRole::Math1) { + } else if (is_state_wg) { DPRINTF0_WG( "Compute[state]: warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, @@ -592,7 +592,7 @@ struct FlatKernelTmaWarpSpecializedKdaFwd { math_barriers, storage.tensors.mainloop); } - } else if (warp_group_role == WarpGroupRole::MathA) { + } else if (is_aux_wg) { DPRINTF0_WG( "Compute[aux]: warp_group_idx:%d, RegisterRequirement:%d\n", warp_group_idx, AuxMmaRegisterRequirement); cutlass::arch::warpgroup_reg_alloc(); diff --git a/csrc/kda/sm90/kernel/options.hpp b/csrc/kda/sm90/kernel/options.hpp index e25fe9d4..cbd93311 100644 --- a/csrc/kda/sm90/kernel/options.hpp +++ b/csrc/kda/sm90/kernel/options.hpp @@ -82,6 +82,9 @@ enum class Tag { kInitStateFromInput, // if true, initialize state by reading global memory instead of zero initialization. kSafeGate, // KDA kElementBetaGmem, // GMEM element type for beta (default float, can be bf16) + kScalarAlpha, // Qwen GDN: one gate value per token/V head, broadcast over D + kStateKVLayout, // Qwen adapter: external state is contiguous [K, V] + kSplitValueDim, // Qwen SM90: split the V=128 feature dimension into two V64 CTAs }; } // namespace kda::sm90::kernel diff --git a/csrc/kda/sm90/kernel/tile_scheduler.hpp b/csrc/kda/sm90/kernel/tile_scheduler.hpp index c70c7a63..4d690d4a 100644 --- a/csrc/kda/sm90/kernel/tile_scheduler.hpp +++ b/csrc/kda/sm90/kernel/tile_scheduler.hpp @@ -27,6 +27,7 @@ struct WorkDesc { int32_t seq_idx; // which sequence to process int32_t qk_head_idx; // head idx for Q/K (the representative of the GVA group) int32_t head_idx; // head idx for V/O/g/beta + int32_t value_tile_idx; // V-feature tile within one value head int64_t tok_offset; // start offset of this sequence in the packed tensor // shape @@ -65,10 +66,12 @@ struct WorkDesc { } }; -// Each block handles a single (seq, v_head) work item; CTAs do not cooperate. +// Each block handles one (seq, v_head, value_tile) work item; CTAs do not cooperate. // GVA optimization: heads_per_group is precomputed on the host and stored in // Params, so the device side does not redo the integer division per CTA. +template struct IndividualTileScheduler { + static_assert(NumValueTiles >= 1); struct Params { dim3 grid; int32_t num_seqs; @@ -93,7 +96,7 @@ struct IndividualTileScheduler { // the integer division. int32_t const heads_per_group = problem_size.num_v_heads / problem_size.num_qk_heads; dim3 grid(0, 1, 1); - grid.x = problem_size.num_seqs * problem_size.num_v_heads; + grid.x = problem_size.num_seqs * problem_size.num_v_heads * NumValueTiles; DPRINTF( "to_underlying_arguments: grid:{.x:%d, .y:%d, .z:%d}, num_seqs:%d, num_qk_heads:%d, num_v_heads:%d, " "heads_per_group:%d\n", @@ -120,8 +123,10 @@ struct IndividualTileScheduler { template CUTE_DEVICE WorkDesc get_next_work(Params params, ProblemSize const& problem_size) { - int32_t seq_idx = blockIdx.x / params.num_v_heads; - int32_t head_idx = blockIdx.x % params.num_v_heads; + int32_t value_tile_idx = blockIdx.x % NumValueTiles; + int32_t work_idx = blockIdx.x / NumValueTiles; + int32_t seq_idx = work_idx / params.num_v_heads; + int32_t head_idx = work_idx % params.num_v_heads; // GVA: use the host-precomputed heads_per_group to avoid device-side division. int32_t qk_head_idx = head_idx / params.heads_per_group; @@ -134,10 +139,12 @@ struct IndividualTileScheduler { } else { scheduled = true; DPRINTF0_W( - "get_next_work: this_work={seq_idx:%d qk_head_idx:%d head_idx:%d tok_offset:%lld seq_len:%lld}\n", + "get_next_work: this_work={seq_idx:%d qk_head_idx:%d head_idx:%d value_tile_idx:%d " + "tok_offset:%lld seq_len:%lld}\n", seq_idx, qk_head_idx, head_idx, + value_tile_idx, s, seq_len); } @@ -146,6 +153,7 @@ struct IndividualTileScheduler { .seq_idx = seq_idx, .qk_head_idx = qk_head_idx, .head_idx = head_idx, + .value_tile_idx = value_tile_idx, .tok_offset = s, .seq_len = seq_len, }; diff --git a/csrc/kda/sm90/prefill_kernel.hpp b/csrc/kda/sm90/prefill_kernel.hpp index 6e54c3a3..94341ee6 100644 --- a/csrc/kda/sm90/prefill_kernel.hpp +++ b/csrc/kda/sm90/prefill_kernel.hpp @@ -51,4 +51,29 @@ launch_kda_fwd_prefill_kernel( int32_t const* raw_cu_seqlens = nullptr, int32_t raw_num_seqs = 0); +// Qwen GDN-only specialization. Alpha is compact chunk-prefix log2 gate +// [packed_tokens, num_v_heads], and state buffers use contiguous [K, V]. +// The generic vector-alpha public API above remains unchanged. +void +launch_qwen35_scalar_kda_fwd_prefill_kernel( + cudaStream_t stream, + void* output, + float* output_state, + void const* q, + void const* k, + void const* v, + float const* input_state, + float const* alpha, + float const* beta, + int32_t const* cu_seqlens, + uint8_t* workspace_buffer, + int32_t num_seqs, + int32_t num_qk_heads, + int32_t num_v_heads, + int32_t head_size, + int64_t total_seqlen, + float scale, + bool has_initial_state, + int32_t sm_count); + } // namespace kda::sm90 diff --git a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh index c53f2ae3..8a482372 100644 --- a/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh +++ b/csrc/kda/sm90/prefill_kernel_kda_fwd_sm90.cuh @@ -38,7 +38,10 @@ template < typename TO, typename TQKV, typename TState, - typename TBeta = float> + typename TBeta = float, + bool ScalarAlpha = false, + bool StateKVLayout = false, + bool SplitValueDim = false> void launch_kda_fwd_prefill_kernel_gbai( cudaStream_t stream, @@ -81,17 +84,21 @@ launch_kda_fwd_prefill_kernel_gbai( using NeedsBetaType = std::conditional_t; using NeedsAlphaType = std::conditional_t; using InitStateType = std::conditional_t; + using Options0 = decltype(add_option(Option{}, DefaultOptions{})); + using Options1 = decltype(add_option(Option{}, Options0{})); + using Options2 = decltype(add_option(Option{}, Options1{})); + using Options3 = decltype(add_option(Option{}, Options2{})); + using Options4 = decltype(add_option(Option{}, Options3{})); + using Options5 = decltype(add_option(Option{}, Options4{})); + using Options6 = decltype(add_option( + Option>{}, + Options5{})); + using Options7 = decltype(add_option( + Option>{}, + Options6{})); using Options = decltype(add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option( - Option{}, - add_option(Option{}, DefaultOptions{}))))))); + Option>{}, + Options7{})); using TileShape = Shape<_64, _64, _128>; using Scheduler = cutlass::gemm::KernelTmaWarpSpecializedCooperative; @@ -137,7 +144,11 @@ launch_kda_fwd_prefill_kernel_gbai( .ptr_K = (T*)k, .dK = {qk_tok_stride, _1{}, head_stride}, .ptr_V = (T*)v, .dV = {v_tok_stride, _1{}, head_stride}, .ptr_O = (T*)output, .dO = {v_tok_stride, _1{}, head_stride}, - .ptr_Alpha = alpha, .dAlpha = {v_tok_stride, _1{}, head_stride}, + .ptr_Alpha = alpha, + .dAlpha = { + ScalarAlpha ? int64_t(num_v_heads) : int64_t(v_tok_stride), + _1{}, + ScalarAlpha ? int32_t(4) : head_stride}, .ptr_output_state = (float*)output_state, .ptr_input_state = (float*)input_state, .scale = scale, diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu index a92481a9..b18b4e80 100644 --- a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill.cu @@ -14,12 +14,22 @@ #include "qwen35_prefill_common.cuh" #include "qwen35_scalar_kda_prefill_kernel.hpp" +#ifdef CULA_SM90A_ENABLED +#include "kda/sm90/prefill_kernel.hpp" +#endif +#ifdef CULA_SM100_ENABLED +#include "kda/sm100/kda_fwd_common.cuh" +#include "qwen35_chunk_state_output_sm100.hpp" +#include "qwen35_chunk_state_output_sm100_ss.hpp" +#endif #include #include #include #include +#include + namespace cula::qwen35::prefill { namespace { @@ -52,6 +62,7 @@ void dispatch_scalar_prefill_for_heads( float* final_state, int batch_size, int seq_len, + int qk_heads, int sequence_count, bool is_varlen, bool has_initial_state) { @@ -70,6 +81,7 @@ void dispatch_scalar_prefill_for_heads( final_state, batch_size, seq_len, + qk_heads, sequence_count, is_varlen, has_initial_state); @@ -92,24 +104,488 @@ void dispatch_scalar_prefill( float* final_state, int batch_size, int seq_len, + int qk_heads, int sequence_count, bool is_varlen, bool has_initial_state) { switch (local_v_heads) { + case 64: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; case 48: - dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 32: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); break; case 24: - dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 16: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); break; case 12: - dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 8: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); break; case 6: - dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, sequence_count, is_varlen, has_initial_state); + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + case 4: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); break; + case 2: + dispatch_scalar_prefill_for_heads(stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); + break; + default: + TORCH_CHECK(false, "unsupported scalar prefill local V-head count: ", local_v_heads); + } +} + +template +void dispatch_scalar_prefill_precomputed_fallback( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const float* g, + const float* beta, + const float* initial_state, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + int local_v_heads, + bool has_initial_state, + const int32_t* unsafe_gate_flags) { +#define CULA_QWEN35_FALLBACK_CASE(HV) \ + case HV: \ + kernel::launch_qwen35_scalar_kda_prefill_precomputed_fallback( \ + stream, q, k, v, g, beta, initial_state, out, final_state, \ + batch_size, seq_len, qk_heads, has_initial_state, unsafe_gate_flags); \ + break + switch (local_v_heads) { + CULA_QWEN35_FALLBACK_CASE(64); + CULA_QWEN35_FALLBACK_CASE(48); + CULA_QWEN35_FALLBACK_CASE(32); + CULA_QWEN35_FALLBACK_CASE(24); + CULA_QWEN35_FALLBACK_CASE(16); + CULA_QWEN35_FALLBACK_CASE(12); + CULA_QWEN35_FALLBACK_CASE(8); + CULA_QWEN35_FALLBACK_CASE(6); + CULA_QWEN35_FALLBACK_CASE(4); + CULA_QWEN35_FALLBACK_CASE(2); + default: + TORCH_CHECK(false, "Unsupported local V head count: ", local_v_heads, "."); } +#undef CULA_QWEN35_FALLBACK_CASE +} + +#ifdef CULA_SM90A_ENABLED +void run_qwen35_chunk_prefill_sm90_bf16( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& a, + const at::Tensor& b, + const at::Tensor& A_log, + const at::Tensor& dt_bias, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + int local_v_heads, + bool has_initial_state, + cudaStream_t stream, + const at::Tensor* precomputed_g = nullptr, + const at::Tensor* precomputed_beta = nullptr) { + const bool use_sm90_chunk = local_v_heads % 4 == 0; + TORCH_CHECK( + (precomputed_g == nullptr) == (precomputed_beta == nullptr), + "precomputed gate and beta must be supplied together."); + constexpr int chunk_size = kernel::kChunkSize; + const int chunks_per_sequence = (seq_len + chunk_size - 1) / chunk_size; + const int total_chunks = batch_size * chunks_per_sequence; + const int64_t total_tokens = static_cast(batch_size) * seq_len; + const auto bf16_options = q.options().dtype(at::kBFloat16); + const auto fp32_options = q.options().dtype(at::kFloat); + const auto int_options = q.options().dtype(at::kInt); + + at::Tensor q_norm = at::empty_like(q, bf16_options); + at::Tensor k_norm = at::empty_like(k, bf16_options); + at::Tensor g = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + at::Tensor g_raw = precomputed_g == nullptr + ? at::empty({batch_size, seq_len, local_v_heads}, fp32_options) + : *precomputed_g; + // Always materialize beta into private workspace. The core ABI input may be + // aliased or reused concurrently and must remain read-only. + at::Tensor beta = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + // HV=6/2 TP shards cannot satisfy Hopper TMA's four-adjacent-head scalar + // gate transaction. Keep the core ABI correct by marking every head for the + // exact recurrent path; divisible-by-four shapes use speculative SM90 KDA. + at::Tensor unsafe_gate_flags = use_sm90_chunk + ? at::zeros({batch_size, local_v_heads}, int_options) + : at::ones({batch_size, local_v_heads}, int_options); + at::Tensor cu_work = at::empty({batch_size + 1}, int_options); + at::Tensor chunk_indices = at::empty({total_chunks, 2}, int_options); + + kernel::launch_qwen35_chunk_preprocess( + stream, + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + precomputed_g == nullptr + ? reinterpret_cast(a.data_ptr()) + : nullptr, + precomputed_g == nullptr + ? reinterpret_cast(b.data_ptr()) + : nullptr, + precomputed_g == nullptr ? A_log.data_ptr() : nullptr, + precomputed_g == nullptr ? dt_bias.data_ptr() : nullptr, + precomputed_g == nullptr ? nullptr : precomputed_g->data_ptr(), + precomputed_beta == nullptr ? nullptr : precomputed_beta->data_ptr(), + precomputed_g != nullptr, + reinterpret_cast<__nv_bfloat16*>(q_norm.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(k_norm.data_ptr()), + g.data_ptr(), + precomputed_g == nullptr ? g_raw.data_ptr() : nullptr, + beta.data_ptr(), + unsafe_gate_flags.data_ptr(), + cu_work.data_ptr(), + chunk_indices.data_ptr(), + batch_size, + seq_len, + qk_heads, + local_v_heads); + + if (use_sm90_chunk) { + const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + at::Tensor workspace = + at::empty({static_cast(sm_count) * 128}, q.options().dtype(at::kByte)); + kda::sm90::launch_qwen35_scalar_kda_fwd_prefill_kernel( + stream, + out.data_ptr(), + final_state.data_ptr(), + q_norm.data_ptr(), + k_norm.data_ptr(), + v.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + g.data_ptr(), + beta.data_ptr(), + cu_work.data_ptr(), + workspace.data_ptr(), + batch_size, + qk_heads, + local_v_heads, + kHeadDimQK, + total_tokens, + rsqrtf(static_cast(kHeadDimQK)), + has_initial_state, + sm_count); + } + + // The SM90 fully-fused safe-gate algebra assumes raw log gates in [-5, 0]. + // Preprocess marks unsafe (sequence, V-head) pairs while it already reads + // the raw gates. A lightweight recurrent launch returns immediately for + // safe heads; unsafe heads overwrite the speculative fast result exactly, + // without a host sync or Python-side branch. + dispatch_scalar_prefill_precomputed_fallback( + stream, + q_norm.data_ptr(), + k_norm.data_ptr(), + v.data_ptr(), + g_raw.data_ptr(), + beta.data_ptr(), + has_initial_state ? initial_state.data_ptr() : nullptr, + out.data_ptr(), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + local_v_heads, + has_initial_state, + unsafe_gate_flags.data_ptr()); +} +#endif + +#ifdef CULA_SM100_ENABLED +// The TS-UMMA implementation is the default SM100 chunk state/output path. +// Keep compile-time escape hatches for A/B comparisons with the WMMA and +// standalone SS prototypes without changing the Python ABI or launch args. +#ifndef CULA_QWEN35_USE_WMMA_CHUNK +#define CULA_QWEN35_USE_WMMA_CHUNK 0 +#endif +#ifndef CULA_QWEN35_USE_TS_CHUNK +#define CULA_QWEN35_USE_TS_CHUNK 1 +#endif + +template +void launch_chunk_state_output_for_heads( + cudaStream_t stream, + const at::Tensor& q_norm, + const at::Tensor& g, + const at::Tensor& Aqk, + const at::Tensor& w, + const at::Tensor& u, + const at::Tensor& kg, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { +#if CULA_QWEN35_USE_WMMA_CHUNK + kernel::launch_qwen35_chunk_state_output( + stream, + reinterpret_cast(q_norm.data_ptr()), + g.data_ptr(), + reinterpret_cast(Aqk.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(u.data_ptr()), + reinterpret_cast(kg.data_ptr()), + has_initial_state ? initial_state.data_ptr() : nullptr, + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + has_initial_state); +#elif CULA_QWEN35_USE_TS_CHUNK + kernel::sm100_ts::launch_qwen35_chunk_state_output_sm100_ts( + stream, + reinterpret_cast(q_norm.data_ptr()), + g.data_ptr(), + reinterpret_cast(Aqk.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(u.data_ptr()), + reinterpret_cast(kg.data_ptr()), + has_initial_state ? initial_state.data_ptr() : nullptr, + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + has_initial_state); +#else + kernel::sm100_ss::launch_qwen35_chunk_state_output_sm100_ss( + stream, + reinterpret_cast(q_norm.data_ptr()), + g.data_ptr(), + reinterpret_cast(Aqk.data_ptr()), + reinterpret_cast(w.data_ptr()), + reinterpret_cast(u.data_ptr()), + reinterpret_cast(kg.data_ptr()), + has_initial_state ? initial_state.data_ptr() : nullptr, + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), + final_state.data_ptr(), + batch_size, + seq_len, + qk_heads, + has_initial_state); +#endif +} + +void launch_chunk_state_output( + int64_t local_v_heads, + cudaStream_t stream, + const at::Tensor& q_norm, + const at::Tensor& g, + const at::Tensor& Aqk, + const at::Tensor& w, + const at::Tensor& u, + const at::Tensor& kg, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { +#define CULA_QWEN35_CHUNK_HEAD_CASE(HV) \ + case HV: \ + launch_chunk_state_output_for_heads( \ + stream, q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, batch_size, seq_len, \ + qk_heads, has_initial_state); \ + break + switch (local_v_heads) { + CULA_QWEN35_CHUNK_HEAD_CASE(64); + CULA_QWEN35_CHUNK_HEAD_CASE(48); + CULA_QWEN35_CHUNK_HEAD_CASE(32); + CULA_QWEN35_CHUNK_HEAD_CASE(24); + CULA_QWEN35_CHUNK_HEAD_CASE(16); + CULA_QWEN35_CHUNK_HEAD_CASE(12); + CULA_QWEN35_CHUNK_HEAD_CASE(8); + CULA_QWEN35_CHUNK_HEAD_CASE(6); + CULA_QWEN35_CHUNK_HEAD_CASE(4); + CULA_QWEN35_CHUNK_HEAD_CASE(2); + default: + TORCH_CHECK(false, "unsupported chunk prefill local V-head count: ", local_v_heads); + } +#undef CULA_QWEN35_CHUNK_HEAD_CASE +} + +void run_qwen35_chunk_prefill_bf16( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& a, + const at::Tensor& b, + const at::Tensor& A_log, + const at::Tensor& dt_bias, + const at::Tensor& initial_state, + const at::Tensor& out, + const at::Tensor& final_state, + int batch_size, + int seq_len, + int qk_heads, + int local_v_heads, + bool has_initial_state, + cudaStream_t stream, + const at::Tensor* precomputed_g = nullptr, + const at::Tensor* precomputed_beta = nullptr) { + TORCH_CHECK( + (precomputed_g == nullptr) == (precomputed_beta == nullptr), + "precomputed gate and beta must be supplied together."); + constexpr int chunk_size = kernel::kChunkSize; + const int chunks_per_sequence = (seq_len + chunk_size - 1) / chunk_size; + const int total_chunks = batch_size * chunks_per_sequence; + const int64_t total_tokens = static_cast(batch_size) * seq_len; + const auto bf16_options = q.options().dtype(at::kBFloat16); + const auto fp32_options = q.options().dtype(at::kFloat); + const auto int_options = q.options().dtype(at::kInt); + + // These tensors are genuine CUDA workspaces; no Python/reference operation + // participates in the numerical result. They are deliberately explicit + // while the chunk path is stabilized, and can later be supplied by an + // inference workspace pool without changing the kernels. + at::Tensor q_norm = at::empty({batch_size, seq_len, qk_heads, kHeadDimQK}, bf16_options); + at::Tensor k_norm = at::empty_like(q_norm); + // Qwen GDN has one scalar gate per token/value-head. Keep it compact and + // route only this adapter through the scalar-G KDA specializations. + at::Tensor g = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + at::Tensor beta = at::empty({batch_size, seq_len, local_v_heads}, fp32_options); + at::Tensor cu_work = at::empty({batch_size + 1}, int_options); + at::Tensor chunk_indices = at::empty({total_chunks, 2}, int_options); + at::Tensor Aqk = at::empty({batch_size, seq_len, local_v_heads, chunk_size}, bf16_options); + at::Tensor Akk = at::empty_like(Aqk); + at::Tensor w = at::empty({batch_size, seq_len, local_v_heads, kHeadDimQK}, bf16_options); + at::Tensor u = at::empty_like(w); + at::Tensor kg = at::empty_like(w); + + kernel::launch_qwen35_chunk_preprocess( + stream, + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + precomputed_g == nullptr + ? reinterpret_cast(a.data_ptr()) + : nullptr, + precomputed_g == nullptr + ? reinterpret_cast(b.data_ptr()) + : nullptr, + precomputed_g == nullptr ? A_log.data_ptr() : nullptr, + precomputed_g == nullptr ? dt_bias.data_ptr() : nullptr, + precomputed_g == nullptr ? nullptr : precomputed_g->data_ptr(), + precomputed_beta == nullptr ? nullptr : precomputed_beta->data_ptr(), + precomputed_g != nullptr && precomputed_beta != nullptr, + reinterpret_cast<__nv_bfloat16*>(q_norm.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(k_norm.data_ptr()), + g.data_ptr(), + nullptr, + beta.data_ptr(), + nullptr, + cu_work.data_ptr(), + chunk_indices.data_ptr(), + batch_size, + seq_len, + qk_heads, + local_v_heads); + + auto* device_prop = at::cuda::getCurrentDeviceProperties(); + const int scheduler_tiles = total_chunks * local_v_heads; + const int scheduler_sms = std::min(device_prop->multiProcessorCount, scheduler_tiles); + KDA_fwd_intra_params intra{}; + intra.total_q_len = static_cast(total_tokens); + intra.b = batch_size; + intra.h_qk = qk_heads; + intra.h_v = local_v_heads; + intra.heads_per_group = local_v_heads / qk_heads; + intra.d = kHeadDimQK; + intra.chunk_size = chunk_size; + intra.scale = rsqrtf(static_cast(kHeadDimQK)); + intra.use_tf32_inverse = false; + intra.unified_gref = true; + intra.is_beta_bf16 = false; + intra.q_ptr = q_norm.data_ptr(); + intra.k_ptr = k_norm.data_ptr(); + intra.g_ptr = g.data_ptr(); + intra.beta_ptr = beta.data_ptr(); + intra.Aqk_out_ptr = Aqk.data_ptr(); + intra.Akk_out_ptr = Akk.data_ptr(); + intra.cu_seqlens_ptr = cu_work.data_ptr(); + intra.chunk_indices_ptr = chunk_indices.data_ptr(); + intra.shape_Akk = cute::make_shape(intra.total_q_len, chunk_size, local_v_heads); + intra.stride_Akk = cute::make_stride(chunk_size * local_v_heads, cute::_1{}, chunk_size); + intra.num_sm = scheduler_sms; + intra.tile_scheduler_params = StaticPersistentTileScheduler::Params{ + total_chunks, + local_v_heads, + intra.heads_per_group, + intra.num_sm, + nullptr}; + kda::sm100::run_kda_fwd_intra_sm100_qwen_scalar_g(intra, stream); + + KDA_fwd_recomp_w_u_params recomp{}; + recomp.total_len = static_cast(total_tokens); + recomp.b = batch_size; + recomp.h_qk = qk_heads; + recomp.h_v = local_v_heads; + recomp.heads_per_group = local_v_heads / qk_heads; + recomp.d = kHeadDimQK; + recomp.chunk_size = chunk_size; + recomp.is_beta_bf16 = false; + recomp.k_ptr = k_norm.data_ptr(); + recomp.v_ptr = v.data_ptr(); + recomp.q_ptr = q_norm.data_ptr(); + recomp.beta_ptr = beta.data_ptr(); + recomp.A_ptr = Akk.data_ptr(); + recomp.g_ptr = g.data_ptr(); + recomp.cu_seqlens_ptr = cu_work.data_ptr(); + recomp.chunk_indices_ptr = chunk_indices.data_ptr(); + recomp.w_out_ptr = w.data_ptr(); + recomp.u_out_ptr = u.data_ptr(); + recomp.kg_out_ptr = kg.data_ptr(); + recomp.qg_out_ptr = nullptr; + recomp.store_qg = false; + recomp.shape_wukg = cute::make_shape(recomp.total_len, kHeadDimQK, local_v_heads); + recomp.stride_wukg = cute::make_stride(kHeadDimQK * local_v_heads, cute::_1{}, kHeadDimQK); + recomp.num_sm = scheduler_sms; + recomp.tile_scheduler_params = StaticPersistentTileScheduler::Params{ + total_chunks, local_v_heads, recomp.heads_per_group, recomp.num_sm, nullptr}; + kda::sm100::run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g(recomp, stream); + + launch_chunk_state_output( + local_v_heads, + stream, + q_norm, + g, + Aqk, + w, + u, + kg, + initial_state, + out, + final_state, + batch_size, + seq_len, + qk_heads, + has_initial_state); } +#endif } // namespace @@ -168,21 +644,28 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { !cu_seqlens.defined() || cu_seqlens.numel() == 0 || cu_seqlens.scalar_type() == at::kInt, "cu_seqlens must be int32 when provided."); - TORCH_CHECK(q.dim() == 4, "q must be [B, T, 48, 128]."); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, "q/k/v must be 4D."); const int64_t B = q.size(0); const int64_t T = q.size(1); - const int64_t local_v_heads = q.size(2); - TORCH_CHECK(decode::is_supported_local_v_heads(static_cast(local_v_heads)), "local V heads must be one of {48, 24, 12, 6}, got ", local_v_heads, "."); + const int64_t qk_heads = q.size(2); + const int64_t local_v_heads = v.size(2); + TORCH_CHECK(qk_heads > 0 && local_v_heads > 0, "q/k/v head counts must be positive."); + TORCH_CHECK(local_v_heads % qk_heads == 0, "local V heads must be divisible by local Q/K heads."); + TORCH_CHECK( + is_supported_scalar_prefill_v_heads(static_cast(local_v_heads)), + "unsupported Qwen scalar prefill local V-head count: ", local_v_heads, "."); TORCH_CHECK( - q.sizes() == at::IntArrayRef({B, T, local_v_heads, kHeadDimQK}), - "q must have shape [B, T, local_v_heads, 128]."); + q.sizes() == at::IntArrayRef({B, T, qk_heads, kHeadDimQK}), + "q must have shape [B, T, local_qk_heads, 128]."); TORCH_CHECK(k.sizes() == q.sizes(), "k must match q shape."); - TORCH_CHECK(v.sizes() == q.sizes(), "v must match q shape."); + TORCH_CHECK( + v.sizes() == at::IntArrayRef({B, T, local_v_heads, kHeadDimV}), + "v must have shape [B, T, local_v_heads, 128]."); TORCH_CHECK(a.dim() == 3 && a.sizes() == at::IntArrayRef({B, T, local_v_heads}), "a must be [B, T, local_v_heads]."); TORCH_CHECK(b.sizes() == a.sizes(), "b must match a shape."); TORCH_CHECK(A_log.dim() == 1 && A_log.size(0) == local_v_heads, "A_log must be [local_v_heads]."); TORCH_CHECK(dt_bias.dim() == 1 && dt_bias.size(0) == local_v_heads, "dt_bias must be [local_v_heads]."); - TORCH_CHECK(out.sizes() == q.sizes(), "out must match q shape."); + TORCH_CHECK(out.sizes() == v.sizes(), "out must match v shape."); const bool is_varlen = cu_seqlens.defined() && cu_seqlens.numel() > 0; const int64_t sequence_count = is_varlen ? cu_seqlens.numel() - 1 : B; @@ -201,7 +684,44 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { } const at::cuda::OptionalCUDAGuard device_guard(device); - cudaStream_t stream = at::cuda::getDefaultCUDAStream(device.index()); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + +#if defined(CULA_SM100_ENABLED) || defined(CULA_SM90A_ENABLED) + // A single packed sequence is equivalent to the fixed-length B=1 layout, + // so it can use the same chunk scheduler without reading cu_seqlens back to + // the host. True multi-sequence varlen remains on the recurrent fallback. + const bool fixed_like_layout = !is_varlen || sequence_count == 1; +#ifdef CULA_SM90A_ENABLED + const bool chunk_head_supported = local_v_heads % 4 == 0; +#else + constexpr bool chunk_head_supported = true; +#endif + if (q.scalar_type() == at::kBFloat16 && T >= 32 && fixed_like_layout && chunk_head_supported) { +#ifdef CULA_SM100_ENABLED + run_qwen35_chunk_prefill_bf16( +#else + run_qwen35_chunk_prefill_sm90_bf16( +#endif + q, + k, + v, + a, + b, + A_log, + dt_bias, + initial_state, + out, + final_state, + static_cast(B), + static_cast(T), + static_cast(qk_heads), + static_cast(local_v_heads), + has_initial_state, + stream); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } +#endif if (q.scalar_type() == at::kHalf) { dispatch_scalar_prefill( @@ -220,6 +740,7 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { final_state.data_ptr(), static_cast(B), static_cast(T), + static_cast(qk_heads), static_cast(sequence_count), is_varlen, has_initial_state); @@ -240,6 +761,7 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { final_state.data_ptr(), static_cast(B), static_cast(T), + static_cast(qk_heads), static_cast(sequence_count), is_varlen, has_initial_state); @@ -247,4 +769,121 @@ void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params) { C10_CUDA_KERNEL_LAUNCH_CHECK(); } +void run_qwen35_scalar_kda_prefill_core(ScalarKdaPrefillCoreParams& params) { +#if !defined(CULA_SM100_ENABLED) && !defined(CULA_SM90A_ENABLED) + TORCH_CHECK(false, "Qwen scalar GDN prefill core requires an SM90 or SM100 build."); +#else + const at::Tensor& q = params.q; + const at::Tensor& k = params.k; + const at::Tensor& v = params.v; + const at::Tensor& gate_raw = params.g; + const at::Tensor& beta_raw = params.beta; + const at::Tensor& initial_state = params.initial_state; + const at::Tensor& cu_seqlens = params.cu_seqlens; + const at::Tensor& out = params.out; + const at::Tensor& final_state = params.final_state; + + TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor."); + const at::Device device = q.device(); + check_tensor_device(k, "k", device); + check_tensor_device(v, "v", device); + check_tensor_device(gate_raw, "g", device); + check_tensor_device(beta_raw, "beta", device); + check_tensor_device(initial_state, "initial_state", device); + check_tensor_device(cu_seqlens, "cu_seqlens", device); + check_tensor_device(out, "out", device); + check_tensor_device(final_state, "final_state", device); + + check_contiguous(q, "q"); + check_contiguous(k, "k"); + check_contiguous(v, "v"); + check_contiguous(gate_raw, "g"); + check_contiguous(beta_raw, "beta"); + check_contiguous(initial_state, "initial_state"); + check_contiguous(cu_seqlens, "cu_seqlens"); + check_contiguous(out, "out"); + check_contiguous(final_state, "final_state"); + + TORCH_CHECK( + q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && + v.scalar_type() == at::kBFloat16 && out.scalar_type() == at::kBFloat16, + "q/k/v/out must be bfloat16 for the chunk core path."); + TORCH_CHECK(gate_raw.scalar_type() == at::kFloat, "g must be float32."); + TORCH_CHECK(beta_raw.scalar_type() == at::kFloat, "beta must be float32."); + TORCH_CHECK(final_state.scalar_type() == at::kFloat, "final_state must be float32."); + TORCH_CHECK( + !initial_state.defined() || initial_state.numel() == 0 || initial_state.scalar_type() == at::kFloat, + "initial_state must be float32 when provided."); + TORCH_CHECK( + !cu_seqlens.defined() || cu_seqlens.numel() == 0 || cu_seqlens.scalar_type() == at::kInt, + "cu_seqlens must be int32 when provided."); + + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, "q/k/v must be 4D."); + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t qk_heads = q.size(2); + const int64_t local_v_heads = v.size(2); + TORCH_CHECK(T >= 32, "the chunk core path requires sequence length >= 32."); + TORCH_CHECK(qk_heads > 0 && local_v_heads > 0, "q/k/v head counts must be positive."); + TORCH_CHECK(local_v_heads % qk_heads == 0, "local V heads must be divisible by local Q/K heads."); + TORCH_CHECK( + is_supported_scalar_prefill_v_heads(static_cast(local_v_heads)), + "unsupported Qwen scalar prefill local V-head count: ", local_v_heads, "."); + TORCH_CHECK( + q.sizes() == at::IntArrayRef({B, T, qk_heads, kHeadDimQK}), + "q must have shape [B, T, local_qk_heads, 128]."); + TORCH_CHECK(k.sizes() == q.sizes(), "k must match q shape."); + TORCH_CHECK( + v.sizes() == at::IntArrayRef({B, T, local_v_heads, kHeadDimV}), + "v must have shape [B, T, local_v_heads, 128]."); + TORCH_CHECK( + gate_raw.sizes() == at::IntArrayRef({B, T, local_v_heads}), + "g must be [B, T, local_v_heads]."); + TORCH_CHECK(beta_raw.sizes() == gate_raw.sizes(), "beta must match g shape."); + TORCH_CHECK(out.sizes() == v.sizes(), "out must match v shape."); + + const bool is_varlen = cu_seqlens.defined() && cu_seqlens.numel() > 0; + const int64_t sequence_count = is_varlen ? cu_seqlens.numel() - 1 : B; + TORCH_CHECK(sequence_count > 0, "sequence_count must be positive."); + TORCH_CHECK(!is_varlen || (B == 1 && sequence_count == 1), + "the chunk core path supports fixed batches or one packed sequence."); + TORCH_CHECK( + final_state.dim() == 4 && + final_state.sizes() == at::IntArrayRef({sequence_count, local_v_heads, kHeadDimQK, kHeadDimV}), + "final_state must be [N, local_v_heads, 128, 128]."); + const bool has_initial_state = initial_state.defined() && initial_state.numel() > 0; + if (has_initial_state) { + TORCH_CHECK(initial_state.sizes() == final_state.sizes(), "initial_state must match final_state shape."); + } + + const at::cuda::OptionalCUDAGuard device_guard(device); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(device.index()); + const at::Tensor empty; +#ifdef CULA_SM100_ENABLED + run_qwen35_chunk_prefill_bf16( +#else + run_qwen35_chunk_prefill_sm90_bf16( +#endif + q, + k, + v, + empty, + empty, + empty, + empty, + initial_state, + out, + final_state, + static_cast(B), + static_cast(T), + static_cast(qk_heads), + static_cast(local_v_heads), + has_initial_state, + stream, + &gate_raw, + &beta_raw); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +#endif +} + } // namespace cula::qwen35::prefill diff --git a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp index c8b2e396..3c60b10d 100644 --- a/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp +++ b/csrc/qwen35/prefill/qwen35_scalar_kda_prefill_kernel.hpp @@ -17,29 +17,40 @@ #include "qwen35_prefill_common.cuh" #include +#include #include +#include namespace cula::qwen35::prefill::kernel { using namespace cute; -template +template struct Qwen35ScalarKdaPrefillKernel { - static constexpr int kThreads = 128; + static constexpr int kWarpSize = 32; + static constexpr int kWarps = kWarpsPerBlock; + static constexpr int kThreads = kWarps * kWarpSize; + static constexpr int kKValuesPerLane = kHeadDimQK / kWarpSize; static constexpr int kHeadDim = kHeadDimQK; - // Keep the scalar CUDA fallback at one V row per CTA for correctness while - // the SM90 chunk/TMA path is being wired in. The previous multi-row V tile - // version exposed a correctness bug with non-zero initial_state; the chunk - // path should own the next parallelization step. - static constexpr int kVTile = 1; + static constexpr int kColumnsPerWarp = 1; + // Each warp owns one independent V column and keeps its complete recurrent + // state in registers. All warps in the CTA reuse one normalized Q/K vector + // and one scalar gate through shared memory. + static constexpr int kVTile = kWarps; static constexpr int kNumVTiles = kHeadDimV / kVTile; static_assert(kHeadDimQK == 128); static_assert(kHeadDimV == 128); + static_assert(kThreads % kWarpSize == 0); + static_assert(kHeadDimQK % kWarpSize == 0); static_assert(kHeadDimV % kVTile == 0); struct SharedStorage { - float scratch[kThreads]; + float q_norm[kHeadDimQK]; + float k_norm[kHeadDimQK]; + float decay; + float beta; + int unsafe_gate; }; static dim3 block_shape() { @@ -66,23 +77,151 @@ struct Qwen35ScalarKdaPrefillKernel { return static_cast(value); } + // A one-warp specialization avoids CTA barriers altogether. It is useful + // for the high-HV/small-T regime where the extra Q/K/gate work is cheaper + // than synchronizing a multi-warp V tile. + CUTE_DEVICE static void run_warp_only( + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ initial_state, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state, + const float* __restrict__ precomputed_g = nullptr, + const float* __restrict__ precomputed_beta = nullptr, + const int32_t* __restrict__ unsafe_gate_flags = nullptr) { + const int lane = static_cast(threadIdx.x) & 31; + int work = static_cast(blockIdx.x); + const int v_row = work % kHeadDimV; + work /= kHeadDimV; + const int hv = work % kLocalVHeads; + const int seq_idx = work / kLocalVHeads; + if (seq_idx >= sequence_count) { + return; + } + const int repeat = kLocalVHeads / qk_heads; + const int qk_h = hv / repeat; + const int token_begin = is_varlen ? static_cast(cu_seqlens[seq_idx]) : seq_idx * seq_len; + const int token_end = is_varlen ? static_cast(cu_seqlens[seq_idx + 1]) : token_begin + seq_len; + const int state_base = ((seq_idx * kLocalVHeads + hv) * kHeadDimQK) * kHeadDimV; + if (precomputed_g != nullptr) { + if (unsafe_gate_flags != nullptr) { + if (unsafe_gate_flags[seq_idx * kLocalVHeads + hv] == 0) { + return; + } + } else { + bool unsafe_gate = false; + for (int token = token_begin + lane; token < token_end; token += kWarpSize) { + const float gate = precomputed_g[token * kLocalVHeads + hv]; + unsafe_gate = unsafe_gate || !isfinite(gate) || gate < -5.0f || gate > 0.0f; + } + if (!__any_sync(0xffffffffu, unsafe_gate)) { + return; + } + } + } + float state_vals[kKValuesPerLane]; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + state_vals[item] = has_initial_state ? initial_state[state_base + kk * kHeadDimV + v_row] : 0.0f; + } + const float scale = rsqrtf(static_cast(kHeadDimQK)); + const float exp_A = precomputed_g == nullptr ? expf(A_log[hv]) : 0.0f; + const float dt = precomputed_g == nullptr ? dt_bias[hv] : 0.0f; + for (int token = token_begin; token < token_end; ++token) { + const int local_t = token - token_begin; + const int qk_base = ((token * qk_heads + qk_h) * kHeadDimQK); + const int v_input_base = ((token * kLocalVHeads + hv) * kHeadDimV); + float q_vals[kKValuesPerLane]; + float k_vals[kKValuesPerLane]; + float q_norm_sq = 0.0f; + float k_norm_sq = 0.0f; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + q_vals[item] = load_as_float(q[qk_base + kk]); + k_vals[item] = load_as_float(k[qk_base + kk]); + q_norm_sq += q_vals[item] * q_vals[item]; + k_norm_sq += k_vals[item] * k_vals[item]; + } + // The SM90 speculative path passes its BF16-normalized Q/K workspace to + // the exact overwrite. Reuse those values verbatim so both paths have + // identical normalization and rounding semantics. + const float q_rnorm = precomputed_g != nullptr + ? scale + : rsqrtf(fmaxf(warp_sum(q_norm_sq), 1.0e-20f)) * scale; + const float k_rnorm = precomputed_g != nullptr + ? 1.0f + : rsqrtf(fmaxf(warp_sum(k_norm_sq), 1.0e-20f)); + float decay = 0.0f; + float beta = 0.0f; + if (lane == 0) { + const int gate_base = token * kLocalVHeads + hv; + if (precomputed_g != nullptr) { + decay = expf(precomputed_g[gate_base]); + beta = precomputed_beta[gate_base]; + } else { + decay = expf(-exp_A * softplus(load_as_float(a[gate_base]) + dt)); + beta = 1.0f / (1.0f + expf(-load_as_float(b[gate_base]))); + } + } + decay = __shfl_sync(0xffffffffu, decay, 0); + beta = __shfl_sync(0xffffffffu, beta, 0); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + k_vals[item] *= k_rnorm; + q_vals[item] *= q_rnorm; + } + float proj_partial = 0.0f; + float out_partial = 0.0f; + const float v_val = __shfl_sync(0xffffffffu, lane == 0 ? load_as_float(v[v_input_base + v_row]) : 0.0f, 0); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + proj_partial += state_vals[item] * k_vals[item]; + } + const float proj = warp_sum(proj_partial); + const float v_new = beta * (v_val - decay * proj); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + state_vals[item] = decay * state_vals[item] + k_vals[item] * v_new; + out_partial += state_vals[item] * q_vals[item]; + } + const float out_acc = warp_sum(out_partial); + if (lane == 0) { + const int out_off = (token * kLocalVHeads + hv) * kHeadDimV + v_row; + out[out_off] = cast_output(out_acc); + } + } +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + final_state[state_base + kk * kHeadDimV + v_row] = state_vals[item]; + } + (void)batch_size; + } + CUTE_DEVICE static float softplus(float x) { return x > 20.0f ? x : log1pf(expf(x)); } - CUTE_DEVICE static float block_sum(float value, SharedStorage& storage, int tid) { - storage.scratch[tid] = value; - __syncthreads(); - - for (int stride = kThreads / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - storage.scratch[tid] += storage.scratch[tid + stride]; - } - __syncthreads(); + CUTE_DEVICE static float warp_sum(float value) { +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); } - const float result = storage.scratch[0]; - __syncthreads(); - return result; + return __shfl_sync(0xffffffffu, value, 0); } CUTE_DEVICE static void run_device( @@ -99,10 +238,21 @@ struct Qwen35ScalarKdaPrefillKernel { float* __restrict__ final_state, int batch_size, int seq_len, + int qk_heads, int sequence_count, bool is_varlen, bool has_initial_state, + const float* __restrict__ precomputed_g, + const float* __restrict__ precomputed_beta, + const int32_t* __restrict__ unsafe_gate_flags, SharedStorage& storage) { + if constexpr (kWarpsPerBlock == 1) { + run_warp_only( + q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, + batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state, + precomputed_g, precomputed_beta, unsafe_gate_flags); + return; + } auto v_work_tiles = make_v_work_tiles(sequence_count); auto work_layout = make_layout(get<1>(v_work_tiles.shape()), LayoutLeft{}); auto work_coord = work_layout.get_hier_coord(static_cast(blockIdx.x)); @@ -111,6 +261,8 @@ struct Qwen35ScalarKdaPrefillKernel { const int seq_idx = static_cast(get<2>(work_coord)); const int v_base = v_tile_idx * kVTile; const int tid = static_cast(threadIdx.x); + const int warp = tid / kWarpSize; + const int lane = tid % kWarpSize; if (hv >= kLocalVHeads || seq_idx >= sequence_count) { return; @@ -119,77 +271,145 @@ struct Qwen35ScalarKdaPrefillKernel { const int token_begin = is_varlen ? static_cast(cu_seqlens[seq_idx]) : seq_idx * seq_len; const int token_end = is_varlen ? static_cast(cu_seqlens[seq_idx + 1]) : token_begin + seq_len; const int state_base = ((seq_idx * kLocalVHeads + hv) * kHeadDimQK) * kHeadDimV; + const int qk_h = hv / (kLocalVHeads / qk_heads); + + // A compact SM90 safe-gate chunk is exact for per-token log gates in + // [-5, 0]. This recurrent fallback is launched after the fast kernel and + // overwrites only heads whose raw gate falls outside that domain. + if (precomputed_g != nullptr) { + if (unsafe_gate_flags != nullptr) { + if (unsafe_gate_flags[seq_idx * kLocalVHeads + hv] == 0) { + return; + } + } else { + if (tid == 0) { + storage.unsafe_gate = 0; + } + __syncthreads(); + for (int token = token_begin + tid; token < token_end; token += kThreads) { + const float gate = precomputed_g[token * kLocalVHeads + hv]; + if (!isfinite(gate) || gate < -5.0f || gate > 0.0f) { + atomicExch(&storage.unsafe_gate, 1); + } + } + __syncthreads(); + if (storage.unsafe_gate == 0) { + return; + } + } + } - const int kk = tid; - float state_vals[kVTile]; + float state_vals[kColumnsPerWarp][kKValuesPerLane]; #pragma unroll - for (int lane = 0; lane < kVTile; ++lane) { - const int v_row = v_base + lane; - const int state_off = state_base + kk * kHeadDimV + v_row; - state_vals[lane] = 0.0f; - if (kk < kHeadDimQK && v_row < kHeadDimV) { - state_vals[lane] = has_initial_state ? initial_state[state_off] : 0.0f; + for (int column = 0; column < kColumnsPerWarp; ++column) { + const int v_row = v_base + warp * kColumnsPerWarp + column; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + const int state_off = state_base + kk * kHeadDimV + v_row; + state_vals[column][item] = has_initial_state ? initial_state[state_off] : 0.0f; } } - __syncthreads(); const float scale = rsqrtf(static_cast(kHeadDimQK)); - const float exp_A = expf(A_log[hv]); - const float dt = dt_bias[hv]; + const float exp_A = precomputed_g == nullptr ? expf(A_log[hv]) : 0.0f; + const float dt = precomputed_g == nullptr ? dt_bias[hv] : 0.0f; for (int token = token_begin; token < token_end; ++token) { - const int local_t = is_varlen ? token : token - token_begin; - const int qkv_base = ((token * kLocalVHeads + hv) * kHeadDimQK); + const int qk_base = ((token * qk_heads + qk_h) * kHeadDimQK); + const int v_base_input = ((token * kLocalVHeads + hv) * kHeadDimV); const int gate_base = token * kLocalVHeads + hv; - const float q_val = kk < kHeadDimQK ? load_as_float(q[qkv_base + kk]) : 0.0f; - const float k_val = kk < kHeadDimQK ? load_as_float(k[qkv_base + kk]) : 0.0f; - const float q_norm_sq = block_sum(q_val * q_val, storage, tid); - const float k_norm_sq = block_sum(k_val * k_val, storage, tid); - const float q_rnorm = rsqrtf(fmaxf(q_norm_sq, 1.0e-20f)) * scale; - const float k_rnorm = rsqrtf(fmaxf(k_norm_sq, 1.0e-20f)); + if (warp == 0) { + float q_vals_raw[kKValuesPerLane]; + float k_vals_raw[kKValuesPerLane]; + float q_norm_sq = 0.0f; + float k_norm_sq = 0.0f; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + q_vals_raw[item] = load_as_float(q[qk_base + kk]); + k_vals_raw[item] = load_as_float(k[qk_base + kk]); + q_norm_sq += q_vals_raw[item] * q_vals_raw[item]; + k_norm_sq += k_vals_raw[item] * k_vals_raw[item]; + } + q_norm_sq = warp_sum(q_norm_sq); + k_norm_sq = warp_sum(k_norm_sq); + const float q_rnorm = precomputed_g != nullptr + ? scale + : rsqrtf(fmaxf(q_norm_sq, 1.0e-20f)) * scale; + const float k_rnorm = precomputed_g != nullptr + ? 1.0f + : rsqrtf(fmaxf(k_norm_sq, 1.0e-20f)); +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + storage.q_norm[kk] = q_vals_raw[item] * q_rnorm; + storage.k_norm[kk] = k_vals_raw[item] * k_rnorm; + } + if (lane == 0) { + if (precomputed_g != nullptr) { + storage.decay = expf(precomputed_g[gate_base]); + storage.beta = precomputed_beta[gate_base]; + } else { + storage.decay = expf(-exp_A * softplus(load_as_float(a[gate_base]) + dt)); + storage.beta = 1.0f / (1.0f + expf(-load_as_float(b[gate_base]))); + } + } + } + __syncthreads(); + + float q_vals[kKValuesPerLane]; + float k_vals[kKValuesPerLane]; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; + q_vals[item] = storage.q_norm[kk]; + k_vals[item] = storage.k_norm[kk]; + } + const float decay = storage.decay; + const float beta = storage.beta; - const float decay = expf(-exp_A * softplus(load_as_float(a[gate_base]) + dt)); - const float beta = 1.0f / (1.0f + expf(-load_as_float(b[gate_base]))); +#pragma unroll + for (int column = 0; column < kColumnsPerWarp; ++column) { + const int v_row = v_base + warp * kColumnsPerWarp + column; + float proj_partial = 0.0f; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + proj_partial += state_vals[column][item] * k_vals[item]; + } + const float proj = warp_sum(proj_partial); - const float k_norm = k_val * k_rnorm; - const float q_norm = q_val * q_rnorm; + float v_val = lane == 0 ? load_as_float(v[v_base_input + v_row]) : 0.0f; + v_val = __shfl_sync(0xffffffffu, v_val, 0); + const float v_new = beta * (v_val - decay * proj); + float out_partial = 0.0f; #pragma unroll - for (int lane = 0; lane < kVTile; ++lane) { - const int v_row = v_base + lane; - if (v_row < kHeadDimV) { - const float proj_partial = kk < kHeadDimQK ? state_vals[lane] * k_norm : 0.0f; - const float proj = block_sum(proj_partial, storage, tid); - - const float v_val = load_as_float(v[qkv_base + v_row]); - const float v_new = beta * (v_val - decay * proj); - - float out_partial = 0.0f; - if (kk < kHeadDimQK) { - const float state_new = decay * state_vals[lane] + k_norm * v_new; - state_vals[lane] = state_new; - out_partial = state_new * q_norm; - } - const float out_acc = block_sum(out_partial, storage, tid); + for (int item = 0; item < kKValuesPerLane; ++item) { + const float state_new = decay * state_vals[column][item] + k_vals[item] * v_new; + state_vals[column][item] = state_new; + out_partial += state_new * q_vals[item]; + } + const float out_acc = warp_sum(out_partial); - if (tid == 0) { - const int out_off = - (((is_varlen ? 0 : seq_idx) * seq_len + local_t) * kLocalVHeads + hv) * kHeadDimV + v_row; - out[out_off] = cast_output(out_acc); - } + if (lane == 0) { + const int out_off = (token * kLocalVHeads + hv) * kHeadDimV + v_row; + out[out_off] = cast_output(out_acc); } } __syncthreads(); } #pragma unroll - for (int lane = 0; lane < kVTile; ++lane) { - const int v_row = v_base + lane; - if (kk < kHeadDimQK && v_row < kHeadDimV) { + for (int column = 0; column < kColumnsPerWarp; ++column) { + const int v_row = v_base + warp * kColumnsPerWarp + column; +#pragma unroll + for (int item = 0; item < kKValuesPerLane; ++item) { + const int kk = lane + item * kWarpSize; const int state_off = state_base + kk * kHeadDimV + v_row; - final_state[state_off] = state_vals[lane]; + final_state[state_off] = state_vals[column][item]; } } @@ -197,7 +417,632 @@ struct Qwen35ScalarKdaPrefillKernel { } }; -template +// The recurrent scalar kernel above is still the lowest-latency path for very +// short prompts. For real prefill lengths, process time in 64-token chunks +// and use tensor cores for the state/output contractions. The preceding +// native-GVA intra/WY stages provide Aqk, W, U, and Kg for this kernel. +static constexpr int kChunkSize = 64; +static constexpr int kValueTile = 64; +static constexpr int kValueTilesPerBlock = kValueTile / 16; +static constexpr int kChunkOutputWarps = 16; +static constexpr int kChunkStateWarps = kChunkOutputWarps; + +struct alignas(128) Qwen35ChunkStateOutputShared { + __nv_bfloat16 state[kHeadDimQK * kValueTile]; + __nv_bfloat16 matrix[kChunkSize * kHeadDimQK]; + __nv_bfloat16 v_new[kChunkSize * kValueTile]; + float gate_exp[kChunkSize]; + float accum[kHeadDimQK * kValueTile]; + __nv_bfloat16 aqk[kChunkSize * kChunkSize]; +}; + +__device__ __forceinline__ float qwen35_bf16_to_float(__nv_bfloat16 value) { + return __bfloat162float(value); +} + +__device__ __forceinline__ __nv_bfloat16 qwen35_float_to_bf16(float value) { + return __float2bfloat16_rn(value); +} + +__device__ __forceinline__ float qwen35_warp_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return __shfl_sync(0xffffffffu, value, 0); +} + +__global__ void qwen35_chunk_qk_norm_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + __nv_bfloat16* __restrict__ q_norm, + __nv_bfloat16* __restrict__ k_norm, + int vector_count) { + const int vector_idx = static_cast(blockIdx.x); + const int lane = static_cast(threadIdx.x); + if (vector_idx >= vector_count) { + return; + } + const int base = vector_idx * kHeadDimQK; + float q_values[4]; + float k_values[4]; + float q_sq = 0.0f; + float k_sq = 0.0f; +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_values[item] = qwen35_bf16_to_float(q[base + kk]); + k_values[item] = qwen35_bf16_to_float(k[base + kk]); + q_sq += q_values[item] * q_values[item]; + k_sq += k_values[item] * k_values[item]; + } + const float q_rnorm = rsqrtf(qwen35_warp_sum(q_sq) + 1.0e-6f); + const float k_rnorm = rsqrtf(qwen35_warp_sum(k_sq) + 1.0e-6f); +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_norm[base + kk] = qwen35_float_to_bf16(q_values[item] * q_rnorm); + k_norm[base + kk] = qwen35_float_to_bf16(k_values[item] * k_rnorm); + } +} + +__device__ __forceinline__ float qwen35_softplus(float x) { + return x > 20.0f ? x : log1pf(expf(x)); +} + +__global__ void qwen35_chunk_gate_kernel( + const __nv_bfloat16* __restrict__ a, + const __nv_bfloat16* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + float* __restrict__ g, + float* __restrict__ beta, + int32_t* __restrict__ cu_seqlens, + int32_t* __restrict__ chunk_indices, + int batch_size, + int seq_len, + int v_heads, + int chunks_per_sequence) { + __shared__ float scan[kChunkSize]; + const int tid = static_cast(threadIdx.x); + int work = static_cast(blockIdx.x); + const int hv = work % v_heads; + work /= v_heads; + const int chunk = work % chunks_per_sequence; + const int seq = work / chunks_per_sequence; + const int local_t = chunk * kChunkSize + tid; + const bool valid = tid < kChunkSize && local_t < seq_len; + const int token = seq * seq_len + local_t; + + if (tid < kChunkSize) { + float log2_decay = 0.0f; + if (valid) { + const int gate_offset = token * v_heads + hv; + const float raw_a = qwen35_bf16_to_float(a[gate_offset]); + const float raw_b = qwen35_bf16_to_float(b[gate_offset]); + const float log_decay = -expf(A_log[hv]) * qwen35_softplus(raw_a + dt_bias[hv]); + log2_decay = log_decay * 1.4426950408889634f; + beta[gate_offset] = 1.0f / (1.0f + expf(-raw_b)); + } + scan[tid] = log2_decay; + } + __syncthreads(); + +#pragma unroll + for (int offset = 1; offset < kChunkSize; offset <<= 1) { + float addend = 0.0f; + if (tid < kChunkSize && tid >= offset) { + addend = scan[tid - offset]; + } + __syncthreads(); + if (tid < kChunkSize) { + scan[tid] += addend; + } + __syncthreads(); + } + + if (tid < kChunkSize) { + const int row_t = chunk * kChunkSize + tid; + if (row_t < seq_len) { + const int row_token = seq * seq_len + row_t; + g[row_token * v_heads + hv] = scan[tid]; + } + } + + if (hv == 0 && tid == 0) { + const int chunk_idx = seq * chunks_per_sequence + chunk; + chunk_indices[chunk_idx * 2] = seq; + chunk_indices[chunk_idx * 2 + 1] = chunk; + if (chunk == 0) { + cu_seqlens[seq] = seq * seq_len; + } + if (chunk == chunks_per_sequence - 1) { + cu_seqlens[seq + 1] = (seq + 1) * seq_len; + } + } +} + +// The Qwen prefill path always needs both normalization and scalar-gate +// preprocessing. Running the two small kernels back-to-back leaves roughly +// 2--3 us of avoidable serialization at short prefill lengths. This fused +// launcher assigns four warps to four Q/K vectors for the first block range, +// then reuses the same 128-thread block shape for the gate scan blocks. The +// two branches are block-uniform, so the gate barriers never involve norm +// threads from another branch and the prefix-sum semantics are unchanged. +#ifndef CULA_QWEN35_FAST_GATE_SCAN +#define CULA_QWEN35_FAST_GATE_SCAN 1 +#endif +#ifndef CULA_QWEN35_PREPROCESS_THREADS +#define CULA_QWEN35_PREPROCESS_THREADS 256 +#endif +#ifndef CULA_QWEN35_PREPROCESS_GATE_FIRST +#define CULA_QWEN35_PREPROCESS_GATE_FIRST 1 +#endif +static_assert(CULA_QWEN35_PREPROCESS_THREADS >= 64, + "Qwen35 preprocess requires at least two warps"); +static_assert(CULA_QWEN35_PREPROCESS_THREADS % 32 == 0, + "Qwen35 preprocess threads must be a multiple of 32"); +template +__global__ void qwen35_chunk_preprocess_fused_kernel( + const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ a, + const __nv_bfloat16* __restrict__ b, + const float* __restrict__ A_log, + const float* __restrict__ dt_bias, + const float* __restrict__ gate_raw, + const float* __restrict__ beta_in, + __nv_bfloat16* __restrict__ q_norm, + __nv_bfloat16* __restrict__ k_norm, + float* __restrict__ g, + float* __restrict__ g_raw_output, + float* __restrict__ beta, + int32_t* __restrict__ unsafe_gate_flags, + int32_t* __restrict__ cu_seqlens, + int32_t* __restrict__ chunk_indices, + int batch_size, + int seq_len, + int qk_heads, + int v_heads, + int vector_count, + int gate_blocks, + int norm_blocks, + int chunks_per_sequence) { + const int block = static_cast(blockIdx.x); +#if CULA_QWEN35_PREPROCESS_GATE_FIRST + if (block >= gate_blocks) { + const int norm_block = block - gate_blocks; +#else + if (block < norm_blocks) { + const int norm_block = block; +#endif + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int vector_idx = norm_block * (static_cast(blockDim.x) / 32) + warp; + if (vector_idx >= vector_count) { + return; + } + const int base = vector_idx * kHeadDimQK; + float q_values[4]; + float k_values[4]; + float q_sq = 0.0f; + float k_sq = 0.0f; +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_values[item] = qwen35_bf16_to_float(q[base + kk]); + k_values[item] = qwen35_bf16_to_float(k[base + kk]); + q_sq += q_values[item] * q_values[item]; + k_sq += k_values[item] * k_values[item]; + } + const float q_rnorm = rsqrtf(qwen35_warp_sum(q_sq) + 1.0e-6f); + const float k_rnorm = rsqrtf(qwen35_warp_sum(k_sq) + 1.0e-6f); +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int kk = lane + item * 32; + q_norm[base + kk] = qwen35_float_to_bf16(q_values[item] * q_rnorm); + k_norm[base + kk] = qwen35_float_to_bf16(k_values[item] * k_rnorm); + } + return; + } + + __shared__ float scan[kChunkSize + 2]; + const int tid = static_cast(threadIdx.x); +#if CULA_QWEN35_PREPROCESS_GATE_FIRST + int work = block; +#else + int work = block - norm_blocks; +#endif + const int hv = work % v_heads; + work /= v_heads; + const int chunk = work % chunks_per_sequence; + const int seq = work / chunks_per_sequence; + const int local_t = chunk * kChunkSize + tid; + const bool valid = tid < kChunkSize && local_t < seq_len; + const int token = seq * seq_len + local_t; + + if (tid < kChunkSize) { + float log2_decay = 0.0f; + if (valid) { + const int gate_offset = token * v_heads + hv; + float raw_log_decay; + if constexpr (UsePrecomputedGate) { + raw_log_decay = gate_raw[gate_offset]; + beta[gate_offset] = beta_in[gate_offset]; + } else { + const float raw_a = qwen35_bf16_to_float(a[gate_offset]); + const float raw_b = qwen35_bf16_to_float(b[gate_offset]); + raw_log_decay = -expf(A_log[hv]) * qwen35_softplus(raw_a + dt_bias[hv]); + beta[gate_offset] = 1.0f / (1.0f + expf(-raw_b)); + } + log2_decay = raw_log_decay * 1.4426950408889634f; + if (g_raw_output != nullptr) { + g_raw_output[gate_offset] = raw_log_decay; + } + if (unsafe_gate_flags != nullptr && + (!isfinite(raw_log_decay) || raw_log_decay < -5.0f || raw_log_decay > 0.0f)) { + atomicExch(&unsafe_gate_flags[seq * v_heads + hv], 1); + } + } + scan[tid] = log2_decay; + } + __syncthreads(); + +#if CULA_QWEN35_FAST_GATE_SCAN + if (tid < kChunkSize) { + const int lane = tid & 31; + const int warp = tid >> 5; + float prefix = scan[tid]; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const float addend = __shfl_up_sync(0xffffffffu, prefix, offset); + if (lane >= offset) { + prefix += addend; + } + } + if (lane == 31) { + scan[kChunkSize + warp] = prefix; + } + scan[tid] = prefix; + } + __syncthreads(); + if (tid >= 32 && tid < kChunkSize) { + scan[tid] += scan[kChunkSize]; + } + __syncthreads(); +#else +#pragma unroll + for (int offset = 1; offset < kChunkSize; offset <<= 1) { + float addend = 0.0f; + if (tid < kChunkSize && tid >= offset) { + addend = scan[tid - offset]; + } + __syncthreads(); + if (tid < kChunkSize) { + scan[tid] += addend; + } + __syncthreads(); + } +#endif + + if (tid < kChunkSize) { + const int row_t = chunk * kChunkSize + tid; + if (row_t < seq_len) { + const int row_token = seq * seq_len + row_t; + g[row_token * v_heads + hv] = scan[tid]; + } + } + + if (hv == 0 && tid == 0) { + const int chunk_idx = seq * chunks_per_sequence + chunk; + chunk_indices[chunk_idx * 2] = seq; + chunk_indices[chunk_idx * 2 + 1] = chunk; + if (chunk == 0) { + cu_seqlens[seq] = seq * seq_len; + } + if (chunk == chunks_per_sequence - 1) { + cu_seqlens[seq + 1] = (seq + 1) * seq_len; + } + } +} + +template +__global__ __launch_bounds__(kChunkStateWarps * 32, 1) void qwen35_chunk_state_output_kernel( + const __nv_bfloat16* __restrict__ q_norm, + const float* __restrict__ g, + const __nv_bfloat16* __restrict__ Aqk, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ u, + const __nv_bfloat16* __restrict__ kg, + const float* __restrict__ initial_state, + __nv_bfloat16* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + using namespace nvcuda; + extern __shared__ char shared_bytes[]; + auto& shared = *reinterpret_cast(shared_bytes); + const int tid = static_cast(threadIdx.x); + const int warp = tid / 32; + int work = static_cast(blockIdx.x); + const int value_tile = work % (kHeadDimV / kValueTile); + work /= (kHeadDimV / kValueTile); + const int hv = work % kLocalVHeads; + const int seq = work / kLocalVHeads; + if (seq >= batch_size) { + return; + } + const int qk_h = hv / (kLocalVHeads / qk_heads); + const int v_base = value_tile * kValueTile; + const int state_global_base = (seq * kLocalVHeads + hv) * kHeadDimQK * kHeadDimV; + + for (int index = tid; index < kHeadDimQK * kValueTile; index += static_cast(blockDim.x)) { + const int kk = index / kValueTile; + const int vv = index % kValueTile; + shared.state[index] = qwen35_float_to_bf16( + has_initial_state ? initial_state[state_global_base + kk * kHeadDimV + v_base + vv] : 0.0f); + } + __syncthreads(); + + const int chunk_count = (seq_len + kChunkSize - 1) / kChunkSize; + const float q_scale = rsqrtf(static_cast(kHeadDimQK)); + for (int chunk = 0; chunk < chunk_count; ++chunk) { + const int chunk_start = chunk * kChunkSize; + const int valid_rows = min(kChunkSize, seq_len - chunk_start); + + for (int row = tid; row < kChunkSize; row += static_cast(blockDim.x)) { + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.gate_exp[row] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[row] = 0.0f; + } + } + __syncthreads(); + for (int index = tid; index < kChunkSize * kHeadDimQK; index += static_cast(blockDim.x)) { + const int row = index / kHeadDimQK; + const int kk = index % kHeadDimQK; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.matrix[index] = w[(token * kLocalVHeads + hv) * kHeadDimQK + kk]; + } else { + shared.matrix[index] = qwen35_float_to_bf16(0.0f); + } + } + __syncthreads(); + + for (int tile = warp; tile < (kChunkSize / 16) * kValueTilesPerBlock; tile += kChunkOutputWarps) { + const int tile_m = tile / kValueTilesPerBlock; + const int tile_n = tile % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment frag_c; + wmma::fill_fragment(frag_c, 0.0f); +#pragma unroll + for (int kk = 0; kk < kHeadDimQK; kk += 16) { + wmma::load_matrix_sync(frag_a, shared.matrix + tile_m * 16 * kHeadDimQK + kk, kHeadDimQK); + wmma::load_matrix_sync(frag_b, shared.state + kk * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(frag_c, frag_a, frag_b, frag_c); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + frag_c, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + + for (int index = tid; index < kChunkSize * kValueTile; index += static_cast(blockDim.x)) { + const int row = index / kValueTile; + const int vv = index % kValueTile; + float value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = qwen35_bf16_to_float(u[(token * kLocalVHeads + hv) * kHeadDimV + v_base + vv]) - + shared.accum[index]; + } + shared.v_new[index] = qwen35_float_to_bf16(value); + } + for (int index = tid; index < kChunkSize * kHeadDimQK; index += static_cast(blockDim.x)) { + const int row = index / kHeadDimQK; + const int kk = index % kHeadDimQK; + float value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = qwen35_bf16_to_float(q_norm[(token * qk_heads + qk_h) * kHeadDimQK + kk]) * + shared.gate_exp[row] * q_scale; + } + shared.matrix[index] = qwen35_float_to_bf16(value); + } + __syncthreads(); + + for (int tile = warp; tile < (kChunkSize / 16) * kValueTilesPerBlock; tile += kChunkOutputWarps) { + const int tile_m = tile / kValueTilesPerBlock; + const int tile_n = tile % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment frag_c; + wmma::fill_fragment(frag_c, 0.0f); +#pragma unroll + for (int kk = 0; kk < kHeadDimQK; kk += 16) { + wmma::load_matrix_sync(frag_a, shared.matrix + tile_m * 16 * kHeadDimQK + kk, kHeadDimQK); + wmma::load_matrix_sync(frag_b, shared.state + kk * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(frag_c, frag_a, frag_b, frag_c); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + frag_c, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + + for (int index = tid; index < kChunkSize * kChunkSize; index += static_cast(blockDim.x)) { + const int row = index / kChunkSize; + const int col = index % kChunkSize; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.aqk[index] = Aqk[(token * kLocalVHeads + hv) * kChunkSize + col]; + } else { + shared.aqk[index] = qwen35_float_to_bf16(0.0f); + } + } + __syncthreads(); + + if (warp < kChunkOutputWarps) { + const int tile_m = warp / kValueTilesPerBlock; + const int tile_n = warp % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment output_acc; + wmma::load_matrix_sync( + output_acc, + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + kValueTile, + wmma::mem_row_major); +#pragma unroll + for (int kk = 0; kk < kChunkSize; kk += 16) { + wmma::load_matrix_sync(frag_a, shared.aqk + tile_m * 16 * kChunkSize + kk, kChunkSize); + wmma::load_matrix_sync(frag_b, shared.v_new + kk * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(output_acc, frag_a, frag_b, output_acc); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + output_acc, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + + for (int index = tid; index < valid_rows * kValueTile; index += static_cast(blockDim.x)) { + const int row = index / kValueTile; + const int vv = index % kValueTile; + const int token = seq * seq_len + chunk_start + row; + out[(token * kLocalVHeads + hv) * kHeadDimV + v_base + vv] = + qwen35_float_to_bf16(shared.accum[index]); + } + for (int index = tid; index < kChunkSize * kHeadDimQK; index += static_cast(blockDim.x)) { + const int row = index / kHeadDimQK; + const int kk = index % kHeadDimQK; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + shared.matrix[index] = kg[(token * kLocalVHeads + hv) * kHeadDimQK + kk]; + } else { + shared.matrix[index] = qwen35_float_to_bf16(0.0f); + } + } + const int last_token = seq * seq_len + chunk_start + valid_rows - 1; + const float chunk_decay = exp2f(g[last_token * kLocalVHeads + hv]); + __syncthreads(); + + for (int tile = warp; tile < (kHeadDimQK / 16) * kValueTilesPerBlock; tile += kChunkStateWarps) { + const int tile_m = tile / kValueTilesPerBlock; + const int tile_n = tile % kValueTilesPerBlock; + wmma::fragment frag_a; + wmma::fragment frag_b; + wmma::fragment frag_c; + wmma::fill_fragment(frag_c, 0.0f); +#pragma unroll + for (int tt = 0; tt < kChunkSize; tt += 16) { + wmma::load_matrix_sync(frag_a, shared.matrix + tt * kHeadDimQK + tile_m * 16, kHeadDimQK); + wmma::load_matrix_sync(frag_b, shared.v_new + tt * kValueTile + tile_n * 16, kValueTile); + wmma::mma_sync(frag_c, frag_a, frag_b, frag_c); + } + wmma::store_matrix_sync( + shared.accum + tile_m * 16 * kValueTile + tile_n * 16, + frag_c, + kValueTile, + wmma::mem_row_major); + } + __syncthreads(); + for (int index = tid; index < kHeadDimQK * kValueTile; index += static_cast(blockDim.x)) { + const float updated = shared.accum[index] + chunk_decay * qwen35_bf16_to_float(shared.state[index]); + shared.state[index] = qwen35_float_to_bf16(updated); + } + __syncthreads(); + } + + for (int index = tid; index < kHeadDimQK * kValueTile; index += static_cast(blockDim.x)) { + const int kk = index / kValueTile; + const int vv = index % kValueTile; + final_state[state_global_base + kk * kHeadDimV + v_base + vv] = qwen35_bf16_to_float(shared.state[index]); + } +} + +inline void launch_qwen35_chunk_preprocess( + cudaStream_t stream, + const __nv_bfloat16* q, + const __nv_bfloat16* k, + const __nv_bfloat16* a, + const __nv_bfloat16* b, + const float* A_log, + const float* dt_bias, + const float* gate_raw, + const float* beta_in, + bool use_precomputed_gate, + __nv_bfloat16* q_norm, + __nv_bfloat16* k_norm, + float* g, + float* g_raw_output, + float* beta, + int32_t* unsafe_gate_flags, + int32_t* cu_seqlens, + int32_t* chunk_indices, + int batch_size, + int seq_len, + int qk_heads, + int v_heads) { + const int vector_count = batch_size * seq_len * qk_heads; + const int chunks = (seq_len + kChunkSize - 1) / kChunkSize; + const int gate_blocks = batch_size * chunks * v_heads; + constexpr int kNormVectorsPerBlock = CULA_QWEN35_PREPROCESS_THREADS / 32; + const int norm_blocks = (vector_count + kNormVectorsPerBlock - 1) / kNormVectorsPerBlock; + if (use_precomputed_gate) { + qwen35_chunk_preprocess_fused_kernel<<< + norm_blocks + gate_blocks, CULA_QWEN35_PREPROCESS_THREADS, 0, stream>>>( + q, k, a, b, A_log, dt_bias, gate_raw, beta_in, + q_norm, k_norm, g, g_raw_output, beta, unsafe_gate_flags, + cu_seqlens, chunk_indices, batch_size, seq_len, qk_heads, v_heads, + vector_count, gate_blocks, norm_blocks, chunks); + } else { + qwen35_chunk_preprocess_fused_kernel<<< + norm_blocks + gate_blocks, CULA_QWEN35_PREPROCESS_THREADS, 0, stream>>>( + q, k, a, b, A_log, dt_bias, gate_raw, beta_in, + q_norm, k_norm, g, g_raw_output, beta, unsafe_gate_flags, + cu_seqlens, chunk_indices, batch_size, seq_len, qk_heads, v_heads, + vector_count, gate_blocks, norm_blocks, chunks); + } +} + +template +inline void launch_qwen35_chunk_state_output( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + auto kernel_fn = &qwen35_chunk_state_output_kernel; + constexpr size_t shared_bytes = sizeof(Qwen35ChunkStateOutputShared); + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes); + const int grid = batch_size * kLocalVHeads * (kHeadDimV / kValueTile); + kernel_fn<<>>( + q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, + batch_size, seq_len, qk_heads, has_initial_state); +} + + +template __global__ void qwen35_scalar_kda_prefill_kernel( const scalar_t* __restrict__ q, const scalar_t* __restrict__ k, @@ -210,13 +1055,17 @@ __global__ void qwen35_scalar_kda_prefill_kernel( const int32_t* __restrict__ cu_seqlens, scalar_t* __restrict__ out, float* __restrict__ final_state, - int batch_size, - int seq_len, - int sequence_count, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, bool is_varlen, - bool has_initial_state) { - __shared__ typename Qwen35ScalarKdaPrefillKernel::SharedStorage storage; - Qwen35ScalarKdaPrefillKernel::run_device( + bool has_initial_state, + const float* __restrict__ precomputed_g, + const float* __restrict__ precomputed_beta, + const int32_t* __restrict__ unsafe_gate_flags) { + __shared__ typename Qwen35ScalarKdaPrefillKernel::SharedStorage storage; + Qwen35ScalarKdaPrefillKernel::run_device( q, k, v, @@ -230,14 +1079,18 @@ __global__ void qwen35_scalar_kda_prefill_kernel( final_state, batch_size, seq_len, + qk_heads, sequence_count, is_varlen, has_initial_state, + precomputed_g, + precomputed_beta, + unsafe_gate_flags, storage); } -template -void launch_qwen35_scalar_kda_prefill_kernel( +template +void launch_qwen35_scalar_kda_prefill_kernel_variant( cudaStream_t stream, const scalar_t* q, const scalar_t* k, @@ -252,12 +1105,14 @@ void launch_qwen35_scalar_kda_prefill_kernel( float* final_state, int batch_size, int seq_len, + int qk_heads, int sequence_count, bool is_varlen, bool has_initial_state) { - const auto grid = Qwen35ScalarKdaPrefillKernel::grid_shape(sequence_count); - const auto block = Qwen35ScalarKdaPrefillKernel::block_shape(); - qwen35_scalar_kda_prefill_kernel<<>>( + using Kernel = Qwen35ScalarKdaPrefillKernel; + const auto grid = Kernel::grid_shape(sequence_count); + const auto block = Kernel::block_shape(); + qwen35_scalar_kda_prefill_kernel<<>>( q, k, v, @@ -271,9 +1126,83 @@ void launch_qwen35_scalar_kda_prefill_kernel( final_state, batch_size, seq_len, + qk_heads, sequence_count, is_varlen, - has_initial_state); + has_initial_state, + nullptr, + nullptr, + nullptr); +} + +template +void launch_qwen35_scalar_kda_prefill_precomputed_fallback( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const float* g, + const float* beta, + const float* initial_state, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state, + const int32_t* unsafe_gate_flags) { + // Safe inputs make this launch an early-return guard. A 16-warp CTA keeps + // that fixed cost to only 8 CTAs per V head, while still providing an exact + // recurrent overwrite for the rare unsafe head. + constexpr int kFallbackWarps = 16; + using Kernel = Qwen35ScalarKdaPrefillKernel; + qwen35_scalar_kda_prefill_kernel + <<>>( + q, + k, + v, + nullptr, + nullptr, + nullptr, + nullptr, + initial_state, + nullptr, + out, + final_state, + batch_size, + seq_len, + qk_heads, + batch_size, + false, + has_initial_state, + g, + beta, + unsafe_gate_flags); +} + +template +void launch_qwen35_scalar_kda_prefill_kernel( + cudaStream_t stream, + const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const float* dt_bias, + const float* initial_state, + const int32_t* cu_seqlens, + scalar_t* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + int sequence_count, + bool is_varlen, + bool has_initial_state) { + launch_qwen35_scalar_kda_prefill_kernel_variant( + stream, q, k, v, a, b, A_log, dt_bias, initial_state, cu_seqlens, out, final_state, + batch_size, seq_len, qk_heads, sequence_count, is_varlen, has_initial_state); } } // namespace cula::qwen35::prefill::kernel diff --git a/tests/test_qwen35_prefill.py b/tests/test_qwen35_prefill.py index 1bca86ff..8f7ab084 100644 --- a/tests/test_qwen35_prefill.py +++ b/tests/test_qwen35_prefill.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math import pathlib import sys @@ -241,6 +242,187 @@ def test_qwen35_scalar_kda_prefill_cuda_supports_local_tp_shards(local_v_heads: torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) +def test_qwen35_scalar_kda_prefill_cuda_long_nonzero_state_and_current_stream(): + """Exercise the tiled CTA path at the main Qwen prefill target length.""" + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(130) + device = torch.device("cuda") + B, T, H, HV, K = 1, 128, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out = torch.empty_like(v) + state = torch.empty_like(initial_state) + empty = torch.empty(0, device=device, dtype=torch.float32) + cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32) + # The extension must honor the current stream; using the default stream + # here would race this event stream in real inference. + stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(stream): + cula_cuda.qwen35_scalar_kda_prefill( + q.contiguous(), + k.contiguous(), + v.contiguous(), + a.contiguous(), + b.contiguous(), + A_log.contiguous(), + dt_bias.contiguous(), + initial_state.contiguous(), + cu_seqlens, + out, + state, + ) + stream.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + +@pytest.mark.parametrize("T", [32, 65, 128, 129]) +def test_qwen35_scalar_kda_prefill_core_cuda_matches_raw_reference(T: int): + """The preprocessed-gate ABI must preserve the raw scalar-kernel result.""" + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core"): + pytest.skip("qwen35_scalar_kda_prefill_core CUDA extension is not available") + + torch.manual_seed(132) + device = torch.device("cuda") + B, H, HV, K = 1, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.randn(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + g = -torch.exp(A_log).view(1, 1, HV) * torch.nn.functional.softplus(a.float() + dt_bias.view(1, 1, HV)) + beta = torch.sigmoid(b.float()) + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out_core, state_core = qwen35_scalar_kda_prefill_core( + q, k, v, g, beta, initial_state=initial_state, backend="cudac" + ) + torch.cuda.synchronize() + torch.testing.assert_close(out_core.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state_core, state_ref, atol=2e-2, rtol=2e-2) + + +def test_qwen35_scalar_kda_prefill_core_sm90_falls_back_for_large_negative_gate(): + """Finite gates below the SM90 safe domain must use the exact fallback.""" + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] != 9 + or cula_cuda is None + or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core") + ): + pytest.skip("SM90 qwen35_scalar_kda_prefill_core CUDA extension is not available") + + torch.manual_seed(177) + device = torch.device("cuda") + B, T, H, HV, K = 1, 32, 16, 48, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.zeros(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = torch.full((HV,), math.log(10.0), device=device, dtype=torch.float32) + dt_bias = torch.zeros(HV, device=device, dtype=torch.float32) + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + g = -torch.exp(A_log).view(1, 1, HV) * torch.nn.functional.softplus(a.float()) + beta = torch.sigmoid(b.float()) + assert g.max().item() < -5.0 + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out_core, state_core = qwen35_scalar_kda_prefill_core( + q, k, v, g, beta, initial_state=initial_state, backend="cudac" + ) + torch.cuda.synchronize() + assert torch.isfinite(out_core).all() + assert torch.isfinite(state_core).all() + torch.testing.assert_close(out_core.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state_core, state_ref, atol=2e-2, rtol=2e-2) + + +def test_qwen35_scalar_kda_prefill_core_sm90_tp8_hv6_exact_fallback(): + """TP8's HV=6 shape cannot use four-head TMA groups but remains correct.""" + if ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] != 9 + or cula_cuda is None + or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core") + ): + pytest.skip("SM90 qwen35_scalar_kda_prefill_core CUDA extension is not available") + + torch.manual_seed(181) + device = torch.device("cuda") + B, T, H, HV, K = 1, 32, 2, 6, 128 + q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(B, T, HV, K, device=device, dtype=torch.bfloat16) + a = torch.zeros(B, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = torch.full((HV,), math.log(0.25), device=device, dtype=torch.float32) + dt_bias = torch.zeros(HV, device=device, dtype=torch.float32) + initial_state = torch.randn(B, HV, K, K, device=device, dtype=torch.float32) * 0.01 + g = -torch.exp(A_log).view(1, 1, HV) * torch.nn.functional.softplus(a.float()) + beta = torch.sigmoid(b.float()) + g_before = g.clone() + beta_before = beta.clone() + + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, backend="reference" + ) + out_core, state_core = qwen35_scalar_kda_prefill_core( + q, k, v, g, beta, initial_state=initial_state, backend="cudac" + ) + torch.cuda.synchronize() + torch.testing.assert_close(out_core.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state_core, state_ref, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(g, g_before, atol=0.0, rtol=0.0) + torch.testing.assert_close(beta, beta_before, atol=0.0, rtol=0.0) + + +def test_qwen35_scalar_kda_prefill_cuda_varlen_multi_sequence(): + if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_scalar_kda_prefill"): + pytest.skip("qwen35_scalar_kda_prefill CUDA extension is not available") + + torch.manual_seed(131) + device = torch.device("cuda") + T, HV, K = 257, 12, 128 + cu_seqlens = torch.tensor([0, 65, 128, T], device=device, dtype=torch.int32) + q = torch.randn(1, T, HV, K, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + a = torch.randn(1, T, HV, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) + dt_bias = torch.randn(HV, device=device, dtype=torch.float32) * 0.1 + initial_state = torch.randn(3, HV, K, K, device=device, dtype=torch.float32) * 0.01 + out_ref, state_ref = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, cu_seqlens=cu_seqlens, backend="reference" + ) + out, state = qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, initial_state=initial_state, cu_seqlens=cu_seqlens, backend="cudac" + ) + torch.cuda.synchronize() + torch.testing.assert_close(out.float(), out_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(state, state_ref, atol=2e-2, rtol=2e-2) + + def test_qwen35_chunk_qk_prefill_sm90_matches_torch(): if not torch.cuda.is_available() or cula_cuda is None or not hasattr(cula_cuda, "qwen35_chunk_qk_prefill_sm90"): import pytest From 875537aa867611ae804d7f1069559a07fe0060ba Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 10:09:21 +0000 Subject: [PATCH 27/35] bench(qwen35): make scalar prefill results auditable --- .../bench_qwen35_scalar_prefill_core.py | 404 ++++++++++++++++++ 1 file changed, 404 insertions(+) create mode 100644 benchmarks/bench_qwen35_scalar_prefill_core.py diff --git a/benchmarks/bench_qwen35_scalar_prefill_core.py b/benchmarks/bench_qwen35_scalar_prefill_core.py new file mode 100644 index 00000000..7b55fe27 --- /dev/null +++ b/benchmarks/bench_qwen35_scalar_prefill_core.py @@ -0,0 +1,404 @@ +#!/usr/bin/env python3 +"""Benchmark the optimized legacy Qwen scalar prefill CUDA kernel. + +The reference path is the real SGLang ``TritonGDNKernel.extend`` path. Both +implementations receive compact native-GVA Q/K and full-HV V/gates. The table +is deliberately a compute-kernel comparison: SGLang's +``fused_gdn_gating`` is warmed and evaluated before timing, while cuLA's +legacy kernel includes its raw ``a/b`` gate conversion, making the comparison +conservative for cuLA. +""" + +from __future__ import annotations + +import argparse +import csv +import importlib.metadata +import json +import pathlib +import statistics +import subprocess +import sys +from typing import Callable + +import torch + +ROOT = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill, qwen35_scalar_kda_prefill_core + + +def _run_text(command: list[str], *, cwd: pathlib.Path) -> str | None: + try: + result = subprocess.run( + command, + cwd=cwd, + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() + + +def _git_head(path: pathlib.Path) -> str: + return _run_text(["git", "rev-parse", "HEAD"], cwd=path) or "unavailable" + + +def _tracked_source_state(path: pathlib.Path) -> str: + status = _run_text( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=path, + ) + if status is None: + return "unavailable" + return "dirty" if status else "clean" + + +def load_shape(path: pathlib.Path, tp: int) -> dict[str, int | torch.dtype | str]: + root = json.loads(path.read_text(encoding="utf-8")) + cfg = root.get("text_config", root) + h_global = int(cfg["linear_num_key_heads"]) + hv_global = int(cfg["linear_num_value_heads"]) + if h_global % tp or hv_global % tp: + raise ValueError(f"TP={tp} must divide global H/HV={h_global}/{hv_global}") + h, hv = h_global // tp, hv_global // tp + if int(cfg["linear_key_head_dim"]) != 128 or int(cfg["linear_value_head_dim"]) != 128: + raise ValueError("This scalar benchmark requires K=V=128") + if hv % h: + raise ValueError(f"local HV={hv} must be divisible by local H={h}") + return { + "model_type": str(cfg.get("model_type", root.get("model_type", "unknown"))), + "h_global": h_global, + "hv_global": hv_global, + "h": h, + "hv": hv, + "dtype": torch.bfloat16, + } + + +def _timed(fn: Callable[[], object], repeats: int) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeats): + fn() + end.record() + end.synchronize() + return start.elapsed_time(end) + + +def _capture_cuda_graph(fn: Callable[[], object]) -> Callable[[], object]: + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + torch.cuda.synchronize() + return graph.replay + + +def _rrms(a: torch.Tensor, b: torch.Tensor) -> float: + af, bf = a.float(), b.float() + return ((af - bf).square().mean().sqrt() / af.square().mean().sqrt().clamp_min(1.0e-8)).item() + + +@torch.inference_mode() +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--sglang-path", type=pathlib.Path, default=pathlib.Path("/sgl-workspace/sglang")) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-lens", type=int, nargs="+", default=(1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096)) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--inner", type=int, default=1, help="kernel calls per CUDA event sample") + parser.add_argument( + "--preheat-iters", + type=int, + default=0, + help="large BF16 GEMMs before timing to stabilize GPU clocks", + ) + parser.add_argument( + "--eager-timing", + action="store_true", + help="time Python launches directly instead of CUDA Graph replay", + ) + parser.add_argument("--random-initial-state", action="store_true") + parser.add_argument("--skip-accuracy", action="store_true") + parser.add_argument( + "--core-only", + action="store_true", + help="compare the preprocessed g/beta calculation core; exclude raw a/b gate conversion on both sides", + ) + parser.add_argument("--csv", type=pathlib.Path, help="write the exact per-shape medians, IQRs, and errors") + parser.add_argument( + "--min-speedup", + type=float, + help="fail unless every acceptance shape reaches this paired-median speedup; requires --core-only", + ) + parser.add_argument( + "--acceptance-seq-lens", + type=int, + nargs="+", + default=(256, 512), + help="sequence lengths checked by --min-speedup", + ) + parser.add_argument( + "--require-clean-source", + action="store_true", + help="fail when tracked source changes are present; ignored build artifacts are allowed", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if args.batch != 1: + raise ValueError("The packed SGLang comparison currently requires --batch 1") + if args.min_speedup is not None and not args.core_only: + parser.error("--min-speedup is an acceptance gate and requires the apples-to-apples --core-only scope") + missing_acceptance_shapes = sorted(set(args.acceptance_seq_lens) - set(args.seq_lens)) + if args.min_speedup is not None and missing_acceptance_shapes: + parser.error(f"acceptance sequence lengths are missing from --seq-lens: {missing_acceptance_shapes}") + source_state = _tracked_source_state(ROOT) + if args.require_clean_source and source_state != "clean": + parser.error(f"formal runs require clean tracked source, got source_state={source_state}") + shape = load_shape(args.config_json, args.tp_size) + import cula.cudac as cula_cuda + + for candidate in (args.sglang_path, args.sglang_path / "python"): + if candidate.exists(): + sys.path.insert(0, str(candidate)) + + from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating + from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel + + device = torch.device("cuda") + sg_kernel = TritonGDNKernel() + if not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core"): + raise RuntimeError("the loaded cuLA extension does not expose qwen35_scalar_kda_prefill_core") + core_op = cula_cuda.qwen35_scalar_kda_prefill_core + extension_module = sys.modules.get(getattr(core_op, "__module__", "")) + extension_path = getattr(extension_module, "__file__", "unavailable") + try: + sglang_version = importlib.metadata.version("sglang") + except importlib.metadata.PackageNotFoundError: + sglang_version = "unavailable" + repo_head = _git_head(ROOT) + sglang_head = _git_head(args.sglang_path) + h, hv = int(shape["h"]), int(shape["hv"]) + print( + f"repo_head={repo_head} tracked_source_state={source_state} " + f"extension={extension_path}" + ) + print( + f"torch={torch.__version__} cuda={torch.version.cuda} " + f"sglang={sglang_version} sglang_head={sglang_head}" + ) + print( + f"device={torch.cuda.get_device_name(device)} config={args.config_json} " + f"model_type={shape['model_type']} TP={args.tp_size} " + f"global H/HV={shape['h_global']}/{shape['hv_global']} local H/HV={h}/{hv}" + ) + print( + f"batch={args.batch} warmup={args.warmup} rep={args.rep} inner={args.inner} " + f"graph={'off' if args.eager_timing else 'on'} " + f"(scope={'preprocessed core' if args.core_only else 'raw CULA gate vs SGLang core'})" + ) + if args.preheat_iters: + heat_a = torch.randn(8192, 8192, device=device, dtype=torch.bfloat16) + heat_b = torch.randn_like(heat_a) + for _ in range(args.preheat_iters): + torch.mm(heat_a, heat_b) + torch.cuda.synchronize() + del heat_a, heat_b + print(f"{'T':>6} {'SGLang ms':>25} {'cuLA ms':>25} {'speedup':>10} {'out rrms':>12} {'state rrms':>12}") + print("-" * 110) + rows: list[dict[str, int | float | str]] = [] + + for seq_len in args.seq_lens: + torch.manual_seed(7000 + seq_len + hv * 17 + args.tp_size) + total = args.batch * seq_len + q = torch.randn(args.batch, seq_len, h, 128, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(args.batch, seq_len, hv, 128, device=device, dtype=torch.bfloat16) + a = torch.randn(args.batch, seq_len, hv, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(hv, device=device, dtype=torch.float32) + dt_bias = torch.randn(hv, device=device, dtype=torch.float32) * 0.1 + state_kv = torch.zeros(args.batch, hv, 128, 128, device=device, dtype=torch.float32) + if args.random_initial_state: + state_kv.normal_(mean=0.0, std=0.01) + state_vk = state_kv.transpose(-1, -2).contiguous() + state_sg = state_vk.clone() + cu = torch.arange(0, total + 1, seq_len, device=device, dtype=torch.int32) + cache_indices = torch.arange(args.batch, device=device, dtype=torch.int32) + + # Keep gating outside both timed calls. This is the same input format + # that SGLang's gdn_backend passes to TritonGDNKernel.extend. + g, beta = fused_gdn_gating(A_log, a.reshape(total, hv), b.reshape(total, hv), dt_bias) + g_core = g.reshape(args.batch, seq_len, hv).contiguous() + beta_core = beta.reshape(args.batch, seq_len, hv).contiguous() + + out_cula = torch.empty_like(v) + state_cula = torch.empty_like(state_kv) + empty_initial = torch.empty(0, device=device, dtype=torch.float32) + + def run_cula() -> None: + # Call the extension ABI directly: output/state allocation and + # Python wrapper overhead are not part of the compute measurement. + cula_state = state_kv + if args.core_only: + cula_cuda.qwen35_scalar_kda_prefill_core( + q, k, v, g_core, beta_core, cula_state, cu, out_cula, state_cula + ) + else: + cula_cuda.qwen35_scalar_kda_prefill( + q, k, v, a, b, A_log, dt_bias, cula_state, cu, out_cula, state_cula + ) + + def reset_sg() -> None: + state_sg.copy_(state_vk) + + def run_sg() -> None: + sg_kernel.extend( + q, + k, + v, + g, + beta, + ssm_states=state_sg, + cache_indices=cache_indices, + query_start_loc=cu, + ) + + # Warm up each path independently, including the first Triton compile. + for _ in range(args.warmup): + run_cula() + reset_sg() + run_sg() + torch.cuda.synchronize() + + out_rrms = float("nan") + state_rrms = float("nan") + if not args.skip_accuracy: + run_cula() + reset_sg() + run_sg() + torch.cuda.synchronize() + # SGLang mutates [V,K], while cuLA writes [K,V]. + reset_sg() + out_sg = sg_kernel.extend( + q, + k, + v, + g, + beta, + ssm_states=state_sg, + cache_indices=cache_indices, + query_start_loc=cu, + )[0] + torch.cuda.synchronize() + out_rrms = _rrms(out_sg, out_cula) + state_rrms = _rrms(state_sg, state_cula.transpose(-1, -2)) + + timed_sg: Callable[[], object] = run_sg + timed_cula: Callable[[], object] = run_cula + if not args.eager_timing: + reset_sg() + timed_sg = _capture_cuda_graph(run_sg) + timed_cula = _capture_cuda_graph(run_cula) + + sg_ms: list[float] = [] + cu_ms: list[float] = [] + # Alternate order and restore the state outside each CUDA event. This + # avoids a systematic clock/thermal bias between the two kernels. + for i in range(args.rep): + if i & 1: + reset_sg() + cu_ms.append(_timed(timed_cula, args.inner) / args.inner) + reset_sg() + sg_ms.append(_timed(timed_sg, args.inner) / args.inner) + else: + reset_sg() + sg_ms.append(_timed(timed_sg, args.inner) / args.inner) + cu_ms.append(_timed(timed_cula, args.inner) / args.inner) + + def middle(xs: list[float]) -> tuple[float, float, float]: + ys = sorted(xs) + return statistics.median(ys), ys[len(ys) // 4], ys[(3 * len(ys)) // 4] + + sg_med, sg_q1, sg_q3 = middle(sg_ms) + cu_med, cu_q1, cu_q3 = middle(cu_ms) + paired = [s / c for s, c in zip(sg_ms, cu_ms)] + speed_med, speed_q1, speed_q3 = middle(paired) + rows.append( + { + "repo_head": repo_head, + "source_state": source_state, + "extension": extension_path, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda or "unavailable", + "sglang_version": sglang_version, + "sglang_head": sglang_head, + "config_json": str(args.config_json.resolve()), + "scope": "core" if args.core_only else "raw_cula_vs_sglang_core", + "cuda_graph": not args.eager_timing, + "random_initial_state": args.random_initial_state, + "warmup": args.warmup, + "rep": args.rep, + "inner": args.inner, + "seq_len": seq_len, + "batch": args.batch, + "tp_size": args.tp_size, + "qk_heads": h, + "v_heads": hv, + "sglang_ms": sg_med, + "sglang_q1_ms": sg_q1, + "sglang_q3_ms": sg_q3, + "cula_ms": cu_med, + "cula_q1_ms": cu_q1, + "cula_q3_ms": cu_q3, + "paired_speedup": speed_med, + "paired_speedup_q1": speed_q1, + "paired_speedup_q3": speed_q3, + "out_rrms": out_rrms, + "state_rrms": state_rrms, + } + ) + print( + f"{seq_len:6d} {sg_med:8.4f} [{sg_q1:8.4f},{sg_q3:8.4f}] " + f"{cu_med:8.4f} [{cu_q1:8.4f},{cu_q3:8.4f}] " + f"{speed_med:8.3f}x [{speed_q1:6.3f},{speed_q3:6.3f}] " + f"{out_rrms:12.3e} {state_rrms:12.3e}" + ) + + if args.csv: + args.csv.parent.mkdir(parents=True, exist_ok=True) + with args.csv.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + if args.min_speedup is not None: + acceptance = {int(row["seq_len"]): float(row["paired_speedup"]) for row in rows} + failures = { + seq_len: acceptance[seq_len] + for seq_len in args.acceptance_seq_lens + if acceptance[seq_len] < args.min_speedup + } + if failures: + formatted = ", ".join(f"T{seq_len}={speedup:.3f}x" for seq_len, speedup in failures.items()) + raise SystemExit(f"speedup acceptance failed (required {args.min_speedup:.3f}x): {formatted}") + print( + "speedup acceptance passed: " + + ", ".join( + f"T{seq_len}={acceptance[seq_len]:.3f}x" for seq_len in args.acceptance_seq_lens + ) + ) + + +if __name__ == "__main__": + main() From 996fbbb882d28f6249595a9ee18959276b29670c Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 10:11:36 +0000 Subject: [PATCH 28/35] bench(qwen35): add scalar prefill NCU target --- benchmarks/profile_qwen35_scalar_prefill.py | 149 ++++++++++++++++++++ 1 file changed, 149 insertions(+) create mode 100644 benchmarks/profile_qwen35_scalar_prefill.py diff --git a/benchmarks/profile_qwen35_scalar_prefill.py b/benchmarks/profile_qwen35_scalar_prefill.py new file mode 100644 index 00000000..390c18f3 --- /dev/null +++ b/benchmarks/profile_qwen35_scalar_prefill.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Nsight Compute target for the legacy Qwen scalar CUDA prefill path. + +This calls only ``qwen35_scalar_kda_prefill_core``. It never resolves or +imports the experimental CuTe prefill backend. Use ``ncu +--profile-from-start off`` so only the launches bracketed by +``cudaProfilerStart/Stop`` are collected. + +Examples:: + + ncu --profile-from-start off --kernel-name-base demangled \ + --kernel-name 'regex:.*qwen35_chunk_state_output_sm100_ts_kernel.*' \ + -o /tmp/qwen35_scalar_t256_state_output \ + python benchmarks/profile_qwen35_scalar_prefill.py \ + --config-json /data/xinhaowei/qwen_configs/Qwen3.5-27B/config.json \ + --seq-len 256 --rep 1 + + ncu --profile-from-start off --kernel-name-base demangled \ + --kernel-name 'regex:.*qwen35_chunk_(preprocess|state_output).*' \ + -o /tmp/qwen35_scalar_t512_stages \ + python benchmarks/profile_qwen35_scalar_prefill.py \ + --config-json /data/xinhaowei/qwen_configs/Qwen3.5-27B/config.json \ + --seq-len 512 --rep 1 +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys + +import torch +import torch.nn.functional as F + +ROOT = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +def load_shape(config_path: pathlib.Path, tp_size: int) -> tuple[int, int]: + root = json.loads(config_path.read_text(encoding="utf-8")) + config = root.get("text_config", root) + global_h = int(config["linear_num_key_heads"]) + global_hv = int(config["linear_num_value_heads"]) + if global_h % tp_size or global_hv % tp_size: + raise ValueError(f"TP={tp_size} must divide global H/HV={global_h}/{global_hv}") + h, hv = global_h // tp_size, global_hv // tp_size + if hv % h: + raise ValueError(f"local HV={hv} must be divisible by local H={h}") + if int(config["linear_key_head_dim"]) != 128 or int(config["linear_value_head_dim"]) != 128: + raise ValueError("the scalar CUDA prefill path requires K=V=128") + return h, hv + + +def extension_path(cula_cuda) -> str: + if not hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core"): + raise RuntimeError("the loaded cuLA extension does not expose qwen35_scalar_kda_prefill_core") + op = cula_cuda.qwen35_scalar_kda_prefill_core + module = sys.modules.get(getattr(op, "__module__", "")) + return str(getattr(module, "__file__", "unavailable")) + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--config-json", type=pathlib.Path, required=True) + parser.add_argument("--tp-size", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--seq-len", type=int, default=256) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--rep", type=int, default=1) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--device", type=int, default=0) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("no CUDA device is available") + if args.rep < 1 or args.warmup < 0: + parser.error("--rep must be positive and --warmup must be non-negative") + if args.seq_len < 32: + parser.error("the chunk scalar CUDA path requires --seq-len >= 32") + + import cula.cudac as cula_cuda + + torch.cuda.set_device(args.device) + device = torch.device("cuda", torch.cuda.current_device()) + h, hv = load_shape(args.config_json, args.tp_size) + torch.manual_seed(args.seed) + + q = torch.randn(1, args.seq_len, h, 128, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, args.seq_len, hv, 128, device=device, dtype=torch.bfloat16) + a = torch.randn(1, args.seq_len, hv, device=device, dtype=torch.bfloat16) + b = torch.randn_like(a) + A_log = -torch.rand(hv, device=device, dtype=torch.float32) + dt_bias = torch.randn(hv, device=device, dtype=torch.float32) * 0.1 + g = -torch.exp(A_log).view(1, 1, hv) * F.softplus(a.float() + dt_bias.view(1, 1, hv)) + beta = torch.sigmoid(b.float()) + initial_state = torch.randn(1, hv, 128, 128, device=device, dtype=torch.float32) * 0.01 + cu_seqlens = torch.tensor([0, args.seq_len], device=device, dtype=torch.int32) + out = torch.empty_like(v) + final_state = torch.empty_like(initial_state) + + def run() -> None: + cula_cuda.qwen35_scalar_kda_prefill_core( + q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + out, + final_state, + ) + + for _ in range(args.warmup): + run() + torch.cuda.synchronize() + + print( + f"profile scalar_cuda_core T={args.seq_len} TP={args.tp_size} H/HV={h}/{hv} " + f"warmup={args.warmup} rep={args.rep} device={torch.cuda.get_device_name(device)}" + ) + print(f"extension={extension_path(cula_cuda)}") + torch.cuda.profiler.start() + for _ in range(args.rep): + run() + torch.cuda.synchronize() + torch.cuda.profiler.stop() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From ad23c1b85ac14c2451b16ae55107e7aaa761a254 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 15:38:30 +0000 Subject: [PATCH 29/35] fix(kda): preserve SM90 context-parallel launcher ABI --- csrc/kda/sm90/kda_fwd_sm90.cu | 10 ++++++++-- csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu | 6 ++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/csrc/kda/sm90/kda_fwd_sm90.cu b/csrc/kda/sm90/kda_fwd_sm90.cu index 48dfaaac..bfa7ff0a 100644 --- a/csrc/kda/sm90/kda_fwd_sm90.cu +++ b/csrc/kda/sm90/kda_fwd_sm90.cu @@ -103,7 +103,10 @@ launch_qwen35_scalar_kda_fwd_prefill_kernel( head_size, total_seqlen, scale, - sm_count); + sm_count, + nullptr, + nullptr, + num_seqs); } else { launch_kda_fwd_prefill_kernel_gbai< true, true, false, true, cutlass::arch::Sm90, bf16, bf16, float, float, true, true, true>( @@ -124,7 +127,10 @@ launch_qwen35_scalar_kda_fwd_prefill_kernel( head_size, total_seqlen, scale, - sm_count); + sm_count, + nullptr, + nullptr, + num_seqs); } } diff --git a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu index f0a51a92..693cc406 100644 --- a/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu +++ b/csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu @@ -86,6 +86,9 @@ launch_kda_fwd_prefill_kernel_gbai< int32_t, int64_t, float, + int32_t, + int32_t const*, + int32_t const*, int32_t); template void @@ -108,6 +111,9 @@ launch_kda_fwd_prefill_kernel_gbai< int32_t, int64_t, float, + int32_t, + int32_t const*, + int32_t const*, int32_t); } // namespace kda::sm90 From 50898ab914e7049befa7e1a1df1c0b03acb52943 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 15:48:24 +0000 Subject: [PATCH 30/35] fix(qwen35): add scalar prefill core parameter definitions --- csrc/qwen35/prefill/qwen35_prefill_common.cuh | 31 +++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/csrc/qwen35/prefill/qwen35_prefill_common.cuh b/csrc/qwen35/prefill/qwen35_prefill_common.cuh index ba3f3b70..ca5c2e16 100644 --- a/csrc/qwen35/prefill/qwen35_prefill_common.cuh +++ b/csrc/qwen35/prefill/qwen35_prefill_common.cuh @@ -41,8 +41,8 @@ struct LayoutPrefillParams { }; struct ScalarKdaPrefillParams { - at::Tensor q; // [B, T, local_v_heads, 128] - at::Tensor k; // [B, T, local_v_heads, 128] + at::Tensor q; // [B, T, local_qk_heads, 128] + at::Tensor k; // [B, T, local_qk_heads, 128] at::Tensor v; // [B, T, local_v_heads, 128] at::Tensor a; // [B, T, local_v_heads] at::Tensor b; // [B, T, local_v_heads] @@ -54,7 +54,34 @@ struct ScalarKdaPrefillParams { at::Tensor final_state; // [N, local_v_heads, 128, 128], float32 }; +// Core-only ABI used for apples-to-apples comparison with SGLang's +// TritonGDNKernel.extend. g/beta are the already materialized per-token +// scalar gate and beta tensors; q/k normalization and chunk-local gate scan +// remain inside the CUDA prefill calculation. +struct ScalarKdaPrefillCoreParams { + at::Tensor q; // [B, T, local_qk_heads, 128], bf16 + at::Tensor k; // [B, T, local_qk_heads, 128], bf16 + at::Tensor v; // [B, T, local_v_heads, 128], bf16 + at::Tensor g; // [B, T, local_v_heads], float32, natural-log gate + at::Tensor beta; // [B, T, local_v_heads], float32 + at::Tensor initial_state; // [N, local_v_heads, 128, 128], float32, may be empty + at::Tensor cu_seqlens; // [N + 1], int32, may be empty + at::Tensor out; // [B, T, local_v_heads, 128], bf16 + at::Tensor final_state; // [N, local_v_heads, 128, 128], float32 +}; + +// All local V-head counts produced by the downloaded Qwen3.5/Qwen3.6 +// configurations at TP={1,2,4,8}. The scalar path accepts compact native-GVA +// Q/K heads and maps each V head to its Q/K group inside the kernel. +inline constexpr bool is_supported_scalar_prefill_v_heads(int local_v_heads) { + return local_v_heads == 64 || local_v_heads == 48 || local_v_heads == 32 || + local_v_heads == 24 || local_v_heads == 16 || local_v_heads == 12 || + local_v_heads == 8 || local_v_heads == 6 || local_v_heads == 4 || + local_v_heads == 2; +} + void run_qwen35_scalar_kda_prefill(ScalarKdaPrefillParams& params); +void run_qwen35_scalar_kda_prefill_core(ScalarKdaPrefillCoreParams& params); void run_qwen35_layout_prefill(LayoutPrefillParams& params); } // namespace cula::qwen35::prefill From db5aaeac1f57022a543efe2a810277e74227462a Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 15:53:12 +0000 Subject: [PATCH 31/35] feat(qwen35): expose scalar prefill core ABI --- csrc/api/pybind.cu | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index 1a04b2f3..c2611f94 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -190,6 +190,31 @@ qwen35_scalar_kda_prefill( cula::qwen35::prefill::run_qwen35_scalar_kda_prefill(params); } +void +qwen35_scalar_kda_prefill_core( + at::Tensor q, + at::Tensor k, + at::Tensor v, + at::Tensor g, + at::Tensor beta, + at::Tensor initial_state, + at::Tensor cu_seqlens, + at::Tensor out, + at::Tensor final_state) { + cula::qwen35::prefill::ScalarKdaPrefillCoreParams params{ + q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + out, + final_state, + }; + cula::qwen35::prefill::run_qwen35_scalar_kda_prefill_core(params); +} + void qwen35_layout_prefill( at::Tensor mixed_qkv_conv, @@ -250,5 +275,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("qwen35_layout_scalar_kda_decode", &qwen35_layout_scalar_kda_decode); m.def("qwen35_layout_prefill", &qwen35_layout_prefill); m.def("qwen35_scalar_kda_prefill", &qwen35_scalar_kda_prefill); + m.def("qwen35_scalar_kda_prefill_core", &qwen35_scalar_kda_prefill_core); m.def("qwen35_chunk_qk_prefill_sm90", &qwen35_chunk_qk_prefill_sm90); } From a08651bb9b98d785956f93e93f6f93193de8c9fd Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 15:57:23 +0000 Subject: [PATCH 32/35] feat(qwen35): connect native-GVA scalar prefill adapter --- cula/ops/qwen35_scalar_kda_prefill.py | 105 +++++++++++++++++++++++--- 1 file changed, 95 insertions(+), 10 deletions(-) diff --git a/cula/ops/qwen35_scalar_kda_prefill.py b/cula/ops/qwen35_scalar_kda_prefill.py index a8ce4204..53d5c486 100644 --- a/cula/ops/qwen35_scalar_kda_prefill.py +++ b/cula/ops/qwen35_scalar_kda_prefill.py @@ -40,12 +40,15 @@ def qwen35_scalar_kda_prefill( """Chunked scalar-gated delta-rule prefill for Qwen3.5.""" if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: - raise ValueError(f"q/k/v must be 4D [B,T,HV,D], got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") - if q.shape != k.shape or q.shape != v.shape: - raise ValueError(f"q/k/v must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") - B, T, HV, K = q.shape - if K != 128 or v.shape[-1] != 128: - raise ValueError(f"Qwen3.5 prefill expects K=V=128, got q={tuple(q.shape)} v={tuple(v.shape)}") + raise ValueError(f"q/k/v must be 4D, got q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}") + if q.shape != k.shape: + raise ValueError(f"q/k must have the same shape, got q={tuple(q.shape)} k={tuple(k.shape)}") + B, T, H, K = q.shape + HV = v.shape[2] + if K != 128 or v.shape[-1] != 128 or v.shape[:2] != q.shape[:2]: + raise ValueError(f"Qwen3.5 prefill expects q/k=[B,T,H,128], v=[B,T,HV,128], got q={tuple(q.shape)} v={tuple(v.shape)}") + if H <= 0 or HV <= 0 or HV % H: + raise ValueError(f"local V heads must be divisible by local Q/K heads, got H={H} HV={HV}") if a.ndim == 2: a = a.unsqueeze(0) if b.ndim == 2: @@ -72,8 +75,9 @@ def qwen35_scalar_kda_prefill( raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_prefill is not available.") if use_cudac: - if HV not in (48, 24, 12, 6): - raise ValueError(f"backend='cudac' supports Qwen3.5 local HV in (48, 24, 12, 6), got {HV}") + supported_hv = (64, 48, 32, 24, 16, 12, 8, 6, 4, 2) + if HV not in supported_hv: + raise ValueError(f"backend='cudac' supports Qwen local HV in {supported_hv}, got {HV}") state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 out = torch.empty_like(v) final_state = torch.empty(state_count, HV, K, K, device=q.device, dtype=torch.float32) @@ -121,13 +125,15 @@ def qwen35_scalar_kda_prefill( dt_bias_f = dt_bias.float() def _run_sequence(batch_idx: int, state_idx: int, start: int, end: int) -> None: + repeat = HV // H for t in range(start, end): for hv in range(HV): + qk_h = hv // repeat state_kv = state[state_idx, hv] decay = torch.exp(-torch.exp(A_log_f[hv]) * torch.nn.functional.softplus(a_f[batch_idx, t, hv] + dt_bias_f[hv])) beta = torch.sigmoid(b_f[batch_idx, t, hv]) - k_vec = k_f[batch_idx, t, hv] - q_vec = q_f[batch_idx, t, hv] + k_vec = k_f[batch_idx, t, qk_h] + q_vec = q_f[batch_idx, t, qk_h] proj = decay * (state_kv.transpose(0, 1) @ k_vec) v_new = beta * (v_f[batch_idx, t, hv] - proj) state_kv_new = decay * state_kv + k_vec.unsqueeze(1) * v_new.unsqueeze(0) @@ -141,3 +147,82 @@ def _run_sequence(batch_idx: int, state_idx: int, start: int, end: int) -> None: for sidx in range(state_count): _run_sequence(0, sidx, int(cu_seqlens[sidx].item()), int(cu_seqlens[sidx + 1].item())) return out, state + + +def qwen35_scalar_kda_prefill_core( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + backend: str = "auto", +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the Qwen GDN calculation after scalar gate/beta preprocessing. + + ``g`` is the natural-log per-token gate before the chunk-local prefix + scan, matching the tensors passed to SGLang's ``TritonGDNKernel.extend``. + The CUDA core still performs Q/K normalization and the prefix scan, while + the raw ``A_log/a/b/dt_bias`` conversion is intentionally outside timing. + """ + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4 or q.shape != k.shape: + raise ValueError("q/k/v must be 4D and q/k must have identical shapes") + B, T, H, K = q.shape + HV = v.shape[2] + if K != 128 or v.shape[:2] != q.shape[:2] or v.shape[-1] != 128 or HV % H: + raise ValueError(f"invalid native-GVA shapes q={tuple(q.shape)} v={tuple(v.shape)}") + if g.ndim == 2: + g = g.unsqueeze(0) + if beta.ndim == 2: + beta = beta.unsqueeze(0) + if g.shape != (B, T, HV) or beta.shape != g.shape: + raise ValueError(f"g/beta must be [B,T,HV], got {tuple(g.shape)} {tuple(beta.shape)}") + if g.dtype != torch.float32 or beta.dtype != torch.float32: + raise ValueError("g and beta must be float32") + if cu_seqlens is not None: + if B != 1 or cu_seqlens.ndim != 1 or cu_seqlens.dtype != torch.int32: + raise ValueError("cu_seqlens must be 1D int32 with B=1") + if initial_state is not None and initial_state.shape[1:] != (HV, K, K): + raise ValueError(f"initial_state must be [N,HV,128,128], got {tuple(initial_state.shape)}") + + use_cudac = ( + backend in ("auto", "cudac") + and cula_cuda is not None + and hasattr(cula_cuda, "qwen35_scalar_kda_prefill_core") + and q.is_cuda + ) + if backend == "cudac" and not use_cudac: + raise RuntimeError("Requested backend='cudac' but qwen35_scalar_kda_prefill_core is unavailable") + if not use_cudac: + raise ValueError("qwen35_scalar_kda_prefill_core currently requires the CUDA backend") + + supported_hv = (64, 48, 32, 24, 16, 12, 8, 6, 4, 2) + if HV not in supported_hv: + raise ValueError(f"backend='cudac' supports local HV in {supported_hv}, got {HV}") + state_count = B if cu_seqlens is None else cu_seqlens.numel() - 1 + out = torch.empty_like(v) + final_state = torch.empty(state_count, HV, K, K, device=q.device, dtype=torch.float32) + initial_state_arg = ( + torch.empty(0, device=q.device, dtype=torch.float32) + if initial_state is None + else initial_state.contiguous() + ) + cu_seqlens_arg = ( + torch.empty(0, device=q.device, dtype=torch.int32) + if cu_seqlens is None + else cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + ) + cula_cuda.qwen35_scalar_kda_prefill_core( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta.contiguous(), + initial_state_arg, + cu_seqlens_arg, + out, + final_state, + ) + return out, final_state From dc327a394afde87f966ec59051d1c90e33ce3170 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 16:00:36 +0000 Subject: [PATCH 33/35] test(qwen35): import scalar prefill core adapter --- tests/test_qwen35_prefill.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_qwen35_prefill.py b/tests/test_qwen35_prefill.py index 8f7ab084..d3279eb3 100644 --- a/tests/test_qwen35_prefill.py +++ b/tests/test_qwen35_prefill.py @@ -25,7 +25,7 @@ from cula.ops.qwen35_conv1d_prefill import qwen35_conv1d_prefill from cula.ops.qwen35_fused_kda_prefill import has_qwen35_fused_kda_prefill, qwen35_fused_kda_prefill from cula.ops.qwen35_layout_prefill import qwen35_layout_prefill, qwen35_layout_prefill_reference -from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill +from cula.ops.qwen35_scalar_kda_prefill import qwen35_scalar_kda_prefill, qwen35_scalar_kda_prefill_core from cula.qwen35.common import Qwen35LinearAttentionConfig from cula.qwen35.runtime import qwen35_linear_attention_prefill From 83c199a4f620f1986a277cc9f2001d32ac981741 Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 16:07:32 +0000 Subject: [PATCH 34/35] feat(qwen35): add SM100 scalar prefill state-output kernels --- .../qwen35_chunk_state_output_sm100.hpp | 639 ++++++++++++++++++ .../qwen35_chunk_state_output_sm100_ss.hpp | 499 ++++++++++++++ 2 files changed, 1138 insertions(+) create mode 100644 csrc/qwen35/prefill/qwen35_chunk_state_output_sm100.hpp create mode 100644 csrc/qwen35/prefill/qwen35_chunk_state_output_sm100_ss.hpp diff --git a/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100.hpp b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100.hpp new file mode 100644 index 00000000..6c2c8d5c --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100.hpp @@ -0,0 +1,639 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +// This header is intentionally independent from +// qwen35_scalar_kda_prefill_kernel.hpp. It is a Blackwell-only prototype for +// replacing that file's WMMA chunk state/output stage without perturbing the +// recurrent fallback or the in-flight scalar-kernel work. + +#if defined(CULA_SM100_ENABLED) + +#include +#include +#include +#include +#include +#include + +#include "kerutils/kerutils.cuh" + +namespace cula::qwen35::prefill::kernel::sm100_ts { + +using namespace cute; + +using bf16 = kerutils::bf16; + +struct alignas(16) Bf16x8 { + bf16 values[8]; +}; + +static constexpr int kHeadDim = 128; +static constexpr int kChunk = 64; +#ifndef CULA_QWEN35_TS_VALUE_TILE +#define CULA_QWEN35_TS_VALUE_TILE 128 +#endif +static constexpr int kValueTile = CULA_QWEN35_TS_VALUE_TILE; +static_assert(kValueTile == 64 || kValueTile == 128); +static constexpr int kTmemThreads = 128; +#ifndef CULA_QWEN35_TS_THREADS +#define CULA_QWEN35_TS_THREADS 352 +#endif +static constexpr int kThreads = CULA_QWEN35_TS_THREADS; + +// TMEM is addressed in 32-bit columns. TS-UMMA requires its M=64 accumulator +// to start at datapath zero, so projection and output use separate 64-column +// regions. (The DP16 packing accepted by the SS recompute mainloop is not a +// legal destination for this TS instruction.) +struct TmemAllocation { + static constexpr uint32_t kStateF32 = 0; // DP 0..31, 128 columns + static constexpr uint32_t kStateBf16 = 128; // DP 0..31, 64 columns + static constexpr uint32_t kResult = 192; // DP 0..31, 64 columns + static constexpr uint32_t kOutput = 256; // DP 0..31, 64 columns + static constexpr int kColumns = 512; +}; + +using SmemLayout64x128K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayout64x64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +// Shape-only layouts used to construct the TMEM A fragments. They are kept +// separate from the 64-row B layouts above because the full-value path uses +// an M=128 TS-UMMA while W/Qg/Aqk still have 64 token rows. +using SmemLayoutValuex128K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutValuex64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +// Logical shape [N=128, K=64]. N/MN-major storage makes both the source KG +// loads (KG is physically [token, K]) and the SMEM stores coalesced while UMMA +// performs the required transpose internally. +using SmemLayout128x64MN = decltype(coalesce( + tile_to_shape( + UMMA::Layout_MN_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using TiledMma64x64K = decltype(make_tiled_mma( + SM100_MMA_F16BF16_TS< + bf16, + bf16, + float, + kValueTile, + kChunk, + UMMA::Major::K, + UMMA::Major::K>{})); + +using TiledMma64x128MN = decltype(make_tiled_mma( + SM100_MMA_F16BF16_TS< + bf16, + bf16, + float, + kValueTile, + kHeadDim, + UMMA::Major::K, + UMMA::Major::MN>{})); + +struct alignas(128) Qwen35ChunkStateOutputSm100Shared { + // Two buffers are required because the two independent TS contractions are + // issued before a single UMMA completion wait. + alignas(128) bf16 operand_b0[kChunk * kHeadDim]; + alignas(128) bf16 operand_b1[kChunk * kHeadDim]; + float gate_exp[kChunk]; + alignas(16) cute::uint64_t mma_barrier; + alignas(16) cute::uint32_t tmem_base_ptr; +}; + +static_assert( + sizeof(Qwen35ChunkStateOutputSm100Shared) <= 48 * 1024, + "SM100 state/output prototype should remain below 48 KiB shared memory"); + +CUTE_DEVICE uint32_t pack_bf16_pair(float x0, float x1) { + union Bf16Bits { + __nv_bfloat16 value; + uint16_t bits; + } lo{}, hi{}; + lo.value = __float2bfloat16_rn(x0); + hi.value = __float2bfloat16_rn(x1); + return static_cast(lo.bits) | (static_cast(hi.bits) << 16); +} + +// Store a [K=128, V=64] FP32 state tile in its transposed TMEM representation +// [V=64, K=128], and create the packed BF16 operand-A shadow at the same time. +CUTE_DEVICE void initialize_state_tmem( + uint32_t tmem_base, + const float* __restrict__ initial_state, + int state_global_base, + int v_base, + bool has_initial_state) { + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + constexpr int kValuesPerWarp = kValueTile / 4; + const bool active = lane < kValuesPerWarp; + const int vv = warp * kValuesPerWarp + (lane & (kValuesPerWarp - 1)); + +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + uint32_t state_bf16[8]; +#pragma unroll + for (int item = 0; item < 16; ++item) { + state_values[item] = active && has_initial_state + ? initial_state[state_global_base + (kk0 + item) * kHeadDim + v_base + vv] + : 0.0f; + } +#pragma unroll + for (int item = 0; item < 8; ++item) { + state_bf16[item] = pack_bf16_pair(state_values[2 * item], state_values[2 * item + 1]); + } + kerutils::tmem_st_32dp32bNx<16>(tmem_base + TmemAllocation::kStateF32 + kk0, state_values); + kerutils::tmem_st_32dp32bNx<8>(tmem_base + TmemAllocation::kStateBf16 + kk0 / 2, state_bf16); + } + cutlass::arch::fence_view_async_tmem_store(); + kerutils::tcgen05_before_thread_sync(); +} + +template +__global__ __launch_bounds__(kThreads, 1) void qwen35_chunk_state_output_sm100_ts_kernel( + const __nv_bfloat16* __restrict__ q_norm, + const float* __restrict__ g, + const __nv_bfloat16* __restrict__ Aqk, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ u, + const __nv_bfloat16* __restrict__ kg, + const float* __restrict__ initial_state, + __nv_bfloat16* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + extern __shared__ char shared_bytes[]; + // CUDA only guarantees the base alignment of an untyped dynamic-shared + // declaration. UMMA SW128 descriptors require 128-byte alignment, so align + // the struct explicitly instead of relying on alignas to move the runtime + // base address. + const uintptr_t shared_addr = reinterpret_cast(shared_bytes); + const uintptr_t shared_aligned_addr = (shared_addr + 127u) & ~uintptr_t(127u); + auto& shared = *reinterpret_cast(shared_aligned_addr); + + const int tid = static_cast(threadIdx.x); + const int lane = tid & 31; + const int warp = tid >> 5; + + int work = static_cast(blockIdx.x); + const int value_tile = work % (kHeadDim / kValueTile); + work /= (kHeadDim / kValueTile); + const int hv = work % kLocalVHeads; + const int seq = work / kLocalVHeads; + if (seq >= batch_size) { + return; + } + + const int heads_per_group = kLocalVHeads / qk_heads; + const int qk_h = hv / heads_per_group; + const int v_base = value_tile * kValueTile; + const int state_global_base = (seq * kLocalVHeads + hv) * kHeadDim * kHeadDim; + + cute::TMEM::Allocator1Sm tmem_allocator{}; + if (warp == 0) { + tmem_allocator.allocate(TmemAllocation::kColumns, &shared.tmem_base_ptr); + // Do not hold the per-SM allocation permit for the whole kernel. + tmem_allocator.release_allocation_lock(); + } + if (tid == 0) { + cute::initialize_barrier(shared.mma_barrier, 1); + } + __syncthreads(); + + const uint32_t tmem_base = shared.tmem_base_ptr; + if (tid < kTmemThreads) { + initialize_state_tmem( + tmem_base, + initial_state, + state_global_base, + v_base, + has_initial_state); + } + __syncthreads(); + kerutils::tcgen05_after_thread_sync(); + + TiledMma64x64K mma_64x64; + TiledMma64x128MN mma_64x128; + + // Construct correctly-shaped TMEM A fragments from fake SMEM tensors. The + // fragment data pointers are then redirected to the explicit TMEM plan. + // A TS operand is a TMEM fragment; the SMEM tensor is shape-only. It must + // use a null SMEM pointer so no real shared-memory base/swizzle offset leaks + // into the generated TMEM fragment layout. + auto fake_state = make_tensor( + make_smem_ptr(static_cast(nullptr)), SmemLayoutValuex128K{}); + auto fake_vnew = make_tensor( + make_smem_ptr(static_cast(nullptr)), SmemLayoutValuex64K{}); + auto t_state_a = mma_64x64.get_slice(_0{}).partition_fragment_A(fake_state); + t_state_a.data() = tmem_base + TmemAllocation::kStateBf16; + auto t_vnew_a_64 = mma_64x64.get_slice(_0{}).partition_fragment_A(fake_vnew); + t_vnew_a_64.data() = tmem_base + TmemAllocation::kStateBf16; + auto t_vnew_a_128 = mma_64x128.get_slice(_0{}).partition_fragment_A(fake_vnew); + t_vnew_a_128.data() = tmem_base + TmemAllocation::kStateBf16; + + auto t_projection = partition_fragment_C( + mma_64x64, Shape, Int>{}); + t_projection.data() = tmem_base + TmemAllocation::kResult; + auto t_output = partition_fragment_C( + mma_64x64, Shape, Int>{}); + t_output.data() = tmem_base + TmemAllocation::kOutput; + auto t_state_acc = partition_fragment_C( + mma_64x128, Shape, Int>{}); + t_state_acc.data() = tmem_base + TmemAllocation::kStateF32; + + int barrier_phase = 0; + const int chunk_count = (seq_len + kChunk - 1) / kChunk; + const float q_scale = rsqrtf(static_cast(kHeadDim)); + const auto* q_norm_bf16 = reinterpret_cast(q_norm); + const auto* Aqk_bf16 = reinterpret_cast(Aqk); + const auto* w_bf16 = reinterpret_cast(w); + const auto* kg_bf16 = reinterpret_cast(kg); + + if constexpr (kPrefetchGate) { + // Seed the first chunk's scalar gate. Later chunks are prefetched while + // the first UMMA pair is in flight, which removes one block barrier from + // every recurrent chunk transition. + if (tid < kChunk) { + if (tid < min(kChunk, seq_len)) { + const int token = seq * seq_len + tid; + shared.gate_exp[tid] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[tid] = 0.0f; + } + } + __syncthreads(); + } + + for (int chunk = 0; chunk < chunk_count; ++chunk) { + const int chunk_start = chunk * kChunk; + const int valid_rows = min(kChunk, seq_len - chunk_start); + + if constexpr (!kPrefetchGate) { + // Keep the short-sequence specialization identical to the lower-latency + // original path; prefetching only pays back after several chunks. + if (tid < kChunk) { + if (tid < valid_rows) { + const int token = seq * seq_len + chunk_start + tid; + shared.gate_exp[tid] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[tid] = 0.0f; + } + } + __syncthreads(); + } + + // First pair: transpose(W @ state) and transpose(Qg @ state). + auto s_w = make_tensor(make_smem_ptr(shared.operand_b0), SmemLayout64x128K{}); + auto s_qg = make_tensor(make_smem_ptr(shared.operand_b1), SmemLayout64x128K{}); + constexpr int kBf16PerVector = 8; + constexpr int kHeadVectors = kHeadDim / kBf16PerVector; + for (int vector_idx = tid; + vector_idx < kChunk * kHeadVectors; + vector_idx += kThreads) { + const int row = vector_idx / kHeadVectors; + const int kk = (vector_idx % kHeadVectors) * kBf16PerVector; + Bf16x8 w_values{}; + Bf16x8 qg_values{}; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + w_values = *reinterpret_cast( + w_bf16 + (token * kLocalVHeads + hv) * kHeadDim + kk); + const Bf16x8 q_values = *reinterpret_cast( + q_norm_bf16 + (token * qk_heads + qk_h) * kHeadDim + kk); +#pragma unroll + for (int item = 0; item < kBf16PerVector; ++item) { + qg_values.values[item] = bf16( + static_cast(q_values.values[item]) * + shared.gate_exp[row] * q_scale); + } + } + *reinterpret_cast(&s_w(row, kk)) = w_values; + *reinterpret_cast(&s_qg(row, kk)) = qg_values; + } + __syncthreads(); + + if (warp == 0) { + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ts(mma_64x64, t_state_a, s_w, t_projection, true); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ts(mma_64x64, t_state_a, s_qg, t_output, true); + cutlass::arch::umma_arrive(&shared.mma_barrier); + } + // Once W/Qg have been staged, gate_exp is dead for this chunk. Use two + // non-issuer warps to prepare the next chunk while UMMA consumes SMEM. + // The epilogue's block barrier below makes these writes visible before + // the next iteration starts loading Qg. + if constexpr (kPrefetchGate) { + if (chunk + 1 < chunk_count && tid >= 32 && tid < 32 + kChunk) { + const int next_row = tid - 32; + const int next_start = chunk_start + kChunk; + const int next_valid_rows = min(kChunk, seq_len - next_start); + if (next_row < next_valid_rows) { + const int token = seq * seq_len + next_start + next_row; + shared.gate_exp[next_row] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[next_row] = 0.0f; + } + } + } + cute::wait_barrier(shared.mma_barrier, barrier_phase); + barrier_phase ^= 1; + + // Load the transposed projection in 16-column slices and create Vnew. + constexpr int kValuesPerWarp = kValueTile / 4; + const bool active_value = lane < kValuesPerWarp; + const int vv = warp * kValuesPerWarp + (lane & (kValuesPerWarp - 1)); + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); +#pragma unroll + for (int row0 = 0; row0 < kChunk; row0 += 16) { + float pair_values[16]; + uint32_t vnew_bf16[8]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kResult + row0, + pair_values); + cutlass::arch::fence_view_async_tmem_load(); +#pragma unroll + for (int item = 0; item < 8; ++item) { + float v0 = 0.0f; + float v1 = 0.0f; + if (active_value) { + const int row_a = row0 + 2 * item; + const int row_b = row_a + 1; + if (row_a < valid_rows) { + const int token = seq * seq_len + chunk_start + row_a; + v0 = __bfloat162float( + u[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv]) - + pair_values[2 * item]; + } + if (row_b < valid_rows) { + const int token = seq * seq_len + chunk_start + row_b; + v1 = __bfloat162float( + u[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv]) - + pair_values[2 * item + 1]; + } + } + vnew_bf16[item] = pack_bf16_pair(v0, v1); + } + kerutils::tmem_st_32dp32bNx<8>( + tmem_base + TmemAllocation::kStateBf16 + row0 / 2, + vnew_bf16); + } + + // Apply the chunk decay to the persistent FP32 state before accumulating + // the KG^T @ Vnew update. DP 16..31 are unused for this tile. + const int last_token = seq * seq_len + chunk_start + valid_rows - 1; + const float chunk_decay = exp2f(g[last_token * kLocalVHeads + hv]); +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + cutlass::arch::fence_view_async_tmem_load(); +#pragma unroll + for (int item = 0; item < 16; ++item) { + state_values[item] *= chunk_decay; + } + kerutils::tmem_st_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + } + cutlass::arch::fence_view_async_tmem_store(); + kerutils::tcgen05_before_thread_sync(); + } + __syncthreads(); + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); + } + + // Second pair: Aqk @ Vnew accumulates into output, while KG^T @ Vnew + // accumulates directly into the decayed FP32 state tile. + auto s_aqk = make_tensor(make_smem_ptr(shared.operand_b0), SmemLayout64x64K{}); + auto s_kg = make_tensor(make_smem_ptr(shared.operand_b1), SmemLayout128x64MN{}); + constexpr int kChunkVectors = kChunk / kBf16PerVector; + for (int vector_idx = tid; + vector_idx < kChunk * kChunkVectors; + vector_idx += kThreads) { + const int row = vector_idx / kChunkVectors; + const int col = (vector_idx % kChunkVectors) * kBf16PerVector; + Bf16x8 values{}; + if (row < valid_rows && col < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + values = *reinterpret_cast( + Aqk_bf16 + (token * kLocalVHeads + hv) * kChunk + col); + } + *reinterpret_cast(&s_aqk(row, col)) = values; + } + // Iterate in KG's physical [token, K] order; MN-major B storage keeps the + // destination N coordinate contiguous too. + for (int vector_idx = tid; + vector_idx < kChunk * kHeadVectors; + vector_idx += kThreads) { + const int row = vector_idx / kHeadVectors; + const int kk = (vector_idx % kHeadVectors) * kBf16PerVector; + Bf16x8 values{}; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + values = *reinterpret_cast( + kg_bf16 + (token * kLocalVHeads + hv) * kHeadDim + kk); + } +#pragma unroll + for (int item = 0; item < kBf16PerVector; ++item) { + s_kg(kk + item, row) = values.values[item]; + } + } + __syncthreads(); + + if (warp == 0) { + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ts(mma_64x64, t_vnew_a_64, s_aqk, t_output, false); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ts(mma_64x128, t_vnew_a_128, s_kg, t_state_acc, false); + cutlass::arch::umma_arrive(&shared.mma_barrier); + } + cute::wait_barrier(shared.mma_barrier, barrier_phase); + barrier_phase ^= 1; + + // Store the completed transposed output. Within each warp, lower-half + // lanes write consecutive V columns for a fixed token. + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); +#pragma unroll + for (int row0 = 0; row0 < kChunk; row0 += 16) { + float pair_values[16]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kOutput + row0, + pair_values); + cutlass::arch::fence_view_async_tmem_load(); + if (active_value) { +#pragma unroll + for (int item = 0; item < 16; ++item) { + const int row = row0 + item; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + out[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv] = + __float2bfloat16_rn(pair_values[item]); + } + } + } + } + + // Refresh the BF16 state shadow for the next chunk. The final chunk also + // writes the exact FP32 persistent state to the public output tensor, but + // does not need a shadow refresh because no later chunk consumes it. + const bool is_last_chunk = chunk + 1 == chunk_count; + if (!is_last_chunk) { +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + uint32_t state_bf16[8]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + cutlass::arch::fence_view_async_tmem_load(); +#pragma unroll + for (int item = 0; item < 8; ++item) { + state_bf16[item] = pack_bf16_pair(state_values[2 * item], state_values[2 * item + 1]); + } + kerutils::tmem_st_32dp32bNx<8>( + tmem_base + TmemAllocation::kStateBf16 + kk0 / 2, + state_bf16); + } + cutlass::arch::fence_view_async_tmem_store(); + kerutils::tcgen05_before_thread_sync(); + } else { +#pragma unroll + for (int kk0 = 0; kk0 < kHeadDim; kk0 += 16) { + float state_values[16]; + kerutils::tmem_ld_32dp32bNx<16>( + tmem_base + TmemAllocation::kStateF32 + kk0, + state_values); + cutlass::arch::fence_view_async_tmem_load(); + if (active_value) { +#pragma unroll + for (int item = 0; item < 16; ++item) { + final_state[state_global_base + (kk0 + item) * kHeadDim + v_base + vv] = + state_values[item]; + } + } + } + } + } + __syncthreads(); + if (tid < kTmemThreads) { + kerutils::tcgen05_after_thread_sync(); + } + } + + __syncthreads(); + if (warp == 0) { + tmem_allocator.free(tmem_base, TmemAllocation::kColumns); + } +} + +template +inline void launch_qwen35_chunk_state_output_sm100_ts_variant( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + auto kernel_fn = + &qwen35_chunk_state_output_sm100_ts_kernel; + constexpr size_t shared_bytes = sizeof(Qwen35ChunkStateOutputSm100Shared) + 127; + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes); + const int grid = batch_size * kLocalVHeads * (kHeadDim / kValueTile); + kernel_fn<<>>( + q_norm, + g, + Aqk, + w, + u, + kg, + initial_state, + out, + final_state, + batch_size, + seq_len, + qk_heads, + has_initial_state); +} + +template +inline void launch_qwen35_chunk_state_output_sm100_ts( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + if (seq_len >= 256) { + launch_qwen35_chunk_state_output_sm100_ts_variant( + stream, q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, + batch_size, seq_len, qk_heads, has_initial_state); + } else { + launch_qwen35_chunk_state_output_sm100_ts_variant( + stream, q_norm, g, Aqk, w, u, kg, initial_state, out, final_state, + batch_size, seq_len, qk_heads, has_initial_state); + } +} + +} // namespace cula::qwen35::prefill::kernel::sm100_ts + +#endif // CULA_SM100_ENABLED diff --git a/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100_ss.hpp b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100_ss.hpp new file mode 100644 index 00000000..a520ad0a --- /dev/null +++ b/csrc/qwen35/prefill/qwen35_chunk_state_output_sm100_ss.hpp @@ -0,0 +1,499 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +// Standalone Blackwell SS-UMMA prototype for the state/output portion of the +// Qwen3.5 scalar prefill path. Keeping this header independent makes it +// possible to compile and inspect the replacement without changing the +// production WMMA kernel or its launcher. + +#if defined(CULA_SM100_ENABLED) + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kerutils/kerutils.cuh" + +namespace cula::qwen35::prefill::kernel::sm100_ss { + +using namespace cute; + +using bf16 = kerutils::bf16; + +static constexpr int kHeadDim = 128; +static constexpr int kChunk = 64; +static constexpr int kValueTile = 64; +static constexpr int kThreads = 128; +// Columns 0..63 hold the two 64-row output accumulators (lower/upper DP +// halves); columns 64..127 hold the independent M128 state update. +static constexpr int kTmemColumns = 128; +static constexpr uint32_t kOutputUpperDp = 16u * 65536u; +static constexpr uint32_t kStateUpdateColumn = 64u; + +// UMMA sees both operands as logical matrices. W/Qg/Aqk/KgT are A +// operands, so K-major is the natural row-major representation. State and +// Vnew are B operands with logical shape [N, K]; MN-major is required here, +// rather than treating their physical [K, V] input representation as an +// ordinary K-major matrix. +using SmemLayoutA64x128K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutStateB64x128MN = decltype(coalesce( + tile_to_shape( + UMMA::Layout_MN_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutA64x64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutVnewB64x64MN = decltype(coalesce( + tile_to_shape( + UMMA::Layout_MN_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +using SmemLayoutKgT128x64K = decltype(coalesce( + tile_to_shape( + UMMA::Layout_K_SW128_Atom{}, + Shape, Int>{}, + Step<_1, _2>{}), + Shape<_1, _1>{})); + +static_assert(cosize_v == kChunk * kHeadDim); +static_assert(cosize_v == kValueTile * kHeadDim); +static_assert(cosize_v == kChunk * kChunk); +static_assert(cosize_v == kValueTile * kChunk); +static_assert(cosize_v == kHeadDim * kChunk); + +using TiledMma64x64 = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS< + bf16, + bf16, + float, + kChunk, + kValueTile, + UMMA::Major::K, + UMMA::Major::MN>{})); + +// State update is Kg^T[M=128,K=64] @ Vnew^T[N=64,K=64]. Keeping M=128 in +// one instruction is important: splitting it into two M64 updates repeats +// the Vnew descriptor traffic and completion synchronization. +using TiledMma128x64 = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS< + bf16, + bf16, + float, + kHeadDim, + kValueTile, + UMMA::Major::K, + UMMA::Major::MN>{})); + +using CompletionPipeline = cutlass::PipelineUmmaAsync<1>; +using CompletionPipelineState = cutlass::PipelineState; +using ClusterShape = Shape<_1, _1, _1>; + +struct alignas(128) Qwen35ChunkStateOutputSm100SsShared { + // Persistent state is laid out as B[N=V,K=head_dim]. + alignas(128) bf16 state[kValueTile * kHeadDim]; + // W/Aqk use this K-major A buffer. It is reused as Aqk after the first + // dual UMMA pair has completed. + alignas(128) bf16 operand_a[kHeadDim * kChunk]; + // Qg/Kg^T use a second K-major A buffer. Keeping the two A operands + // separate allows each pair of contractions to share one completion wait. + alignas(128) bf16 operand_a_aux[kHeadDim * kChunk]; + // Vnew is the MN-major B operand for both Aqk and the M128 state update. + alignas(128) bf16 vnew[kValueTile * kChunk]; + float gate_exp[kChunk]; + alignas(16) uint32_t tmem_base_ptr; + alignas(16) typename CompletionPipeline::SharedStorage completion; +}; + +static_assert( + sizeof(Qwen35ChunkStateOutputSm100SsShared) <= 64 * 1024, + "SS-UMMA state/output prototype must remain below 64 KiB shared memory"); + +CUTE_DEVICE void release_ss_mma_result( + CompletionPipeline& completion, + CompletionPipelineState& consumer_state) { + // TMEM loads performed by the epilogue must become visible before the + // consumer marks the single pipeline stage reusable. + kerutils::tcgen05_before_thread_sync(); + completion.consumer_release(consumer_state); + ++consumer_state; + __syncthreads(); +} + +template +__global__ __launch_bounds__(kThreads, 1) void qwen35_chunk_state_output_sm100_ss_kernel( + const __nv_bfloat16* __restrict__ q_norm, + const float* __restrict__ g, + const __nv_bfloat16* __restrict__ Aqk, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ u, + const __nv_bfloat16* __restrict__ kg, + const float* __restrict__ initial_state, + __nv_bfloat16* __restrict__ out, + float* __restrict__ final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + extern __shared__ char shared_bytes[]; + auto& shared = *reinterpret_cast(shared_bytes); + + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + + int work = static_cast(blockIdx.x); + const int value_tile = work % (kHeadDim / kValueTile); + work /= (kHeadDim / kValueTile); + const int hv = work % kLocalVHeads; + const int seq = work / kLocalVHeads; + if (seq >= batch_size) { + return; + } + + const int qk_h = hv / (kLocalVHeads / qk_heads); + const int v_base = value_tile * kValueTile; + const int state_global_base = (seq * kLocalVHeads + hv) * kHeadDim * kHeadDim; + + auto s_state = make_tensor( + make_smem_ptr(shared.state), SmemLayoutStateB64x128MN{}); + auto s_a_64x128 = make_tensor( + make_smem_ptr(shared.operand_a), SmemLayoutA64x128K{}); + auto s_a_qg = make_tensor( + make_smem_ptr(shared.operand_a_aux), SmemLayoutA64x128K{}); + auto s_a_aqk = make_tensor( + make_smem_ptr(shared.operand_a), SmemLayoutA64x64K{}); + auto s_a_kg = make_tensor( + make_smem_ptr(shared.operand_a_aux), SmemLayoutKgT128x64K{}); + auto s_a_128x64 = make_tensor( + make_smem_ptr(shared.operand_a), SmemLayoutKgT128x64K{}); + auto s_vnew = make_tensor( + make_smem_ptr(shared.vnew), SmemLayoutVnewB64x64MN{}); + + for (int index = tid; index < kValueTile * kHeadDim; index += kThreads) { + const int vv = index / kHeadDim; + const int kk = index % kHeadDim; + const float value = has_initial_state + ? initial_state[state_global_base + kk * kHeadDim + v_base + vv] + : 0.0f; + s_state(vv, kk) = bf16(value); + } + + // One completion pipeline is sufficient because the four contractions are + // issued as two dual-UMMA pairs. Warp 0 both issues UMMA and participates + // in the 128-thread epilogue. + typename CompletionPipeline::Params completion_params; + completion_params.producer_arv_count = 1; + completion_params.consumer_arv_count = kThreads; + completion_params.initializing_warp = 0; + completion_params.role = warp == 0 + ? CompletionPipeline::ThreadCategory::ProducerConsumer + : CompletionPipeline::ThreadCategory::Consumer; + CompletionPipeline completion( + shared.completion, completion_params, ClusterShape{}); + + cute::TMEM::Allocator1Sm tmem_allocator{}; + if (warp == 0) { + tmem_allocator.allocate(kTmemColumns, &shared.tmem_base_ptr); + tmem_allocator.release_allocation_lock(); + } + __syncthreads(); + + TiledMma64x64 mma_64x64; + TiledMma128x64 mma_128x64; + auto t_acc_64x64_lower = partition_fragment_C( + mma_64x64, Shape, Int>{}); + auto t_acc_64x64_upper = partition_fragment_C( + mma_64x64, Shape, Int>{}); + auto t_acc_128x64_state = partition_fragment_C( + mma_128x64, Shape, Int>{}); + t_acc_64x64_lower.data() = shared.tmem_base_ptr; + t_acc_64x64_upper.data() = shared.tmem_base_ptr + kOutputUpperDp; + t_acc_128x64_state.data() = + shared.tmem_base_ptr + kStateUpdateColumn; + + auto c64 = make_identity_tensor( + Shape, Int>{}); + auto c128 = make_identity_tensor( + Shape, Int>{}); + auto t_c64 = mma_64x64.get_slice(_0{}).partition_C(c64); + auto t_c128 = mma_128x64.get_slice(_0{}).partition_C(c128); + + CompletionPipelineState producer_state = + cutlass::make_producer_start_state(); + CompletionPipelineState consumer_state; + + const auto* q_bf16 = reinterpret_cast(q_norm); + const auto* aqk_bf16 = reinterpret_cast(Aqk); + const auto* w_bf16 = reinterpret_cast(w); + const auto* u_bf16 = reinterpret_cast(u); + const auto* kg_bf16 = reinterpret_cast(kg); + const int chunk_count = (seq_len + kChunk - 1) / kChunk; + const float q_scale = rsqrtf(static_cast(kHeadDim)); + + for (int chunk = 0; chunk < chunk_count; ++chunk) { + const int chunk_start = chunk * kChunk; + const int valid_rows = min(kChunk, seq_len - chunk_start); + + if (tid < kChunk) { + if (tid < valid_rows) { + const int token = seq * seq_len + chunk_start + tid; + shared.gate_exp[tid] = exp2f(g[token * kLocalVHeads + hv]); + } else { + shared.gate_exp[tid] = 0.0f; + } + } + + // Pair 1: projection = W[64,128] @ state^T[128,64] in the lower DP + // half, and output_base = Qg[64,128] @ state^T in the upper DP half. + // Both A operands are staged before the single UMMA completion signal. + for (int index = tid; index < kChunk * kHeadDim; index += kThreads) { + const int row = index / kHeadDim; + const int kk = index % kHeadDim; + bf16 value = bf16(0.0f); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = w_bf16[(token * kLocalVHeads + hv) * kHeadDim + kk]; + } + s_a_64x128(row, kk) = value; + float q_value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + q_value = static_cast( + q_bf16[(token * qk_heads + qk_h) * kHeadDim + kk]) * + shared.gate_exp[row] * q_scale; + } + s_a_qg(row, kk) = bf16(q_value); + } + __syncthreads(); + + if (warp == 0) { + completion.producer_acquire(producer_state); + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ss( + mma_64x64, + s_a_64x128, + s_state, + t_acc_64x64_lower, + true); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ss( + mma_64x64, + s_a_qg, + s_state, + t_acc_64x64_upper, + true); + completion.producer_commit(producer_state); + ++producer_state; + } + completion.consumer_wait(consumer_state); + kerutils::tcgen05_after_thread_sync(); + + { + auto tiled_t2r = make_tmem_copy( + SM100_TMEM_LOAD_16dp256b8x{}, t_acc_64x64_lower); + auto thr_t2r = tiled_t2r.get_slice(tid); + auto t_src = thr_t2r.partition_S(t_acc_64x64_lower); + auto t_coord = thr_t2r.partition_D(t_c64); + auto r_acc = make_tensor(shape(t_coord)); + copy(tiled_t2r, t_src, r_acc); + cutlass::arch::fence_view_async_tmem_load(); + CUTE_UNROLL + for (int item = 0; item < size(r_acc); ++item) { + const auto coord = t_coord(item); + const int row = static_cast(get<0>(coord)); + const int vv = static_cast(get<1>(coord)); + float value = 0.0f; + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = static_cast( + u_bf16[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv]) - + r_acc(item); + } + s_vnew(vv, row) = bf16(value); + } + } + cutlass::arch::fence_view_async_shared(); + release_ss_mma_result(completion, consumer_state); + + // Pair 2: output += Aqk[64,64] @ Vnew^T in the upper DP half, while the + // state update Kg^T[128,64] @ Vnew^T is accumulated in a separate M128 + // TMEM fragment. The two independent UMMAs share one completion wait. + for (int index = tid; index < kChunk * kChunk; index += kThreads) { + const int row = index / kChunk; + const int col = index % kChunk; + bf16 value = bf16(0.0f); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = aqk_bf16[(token * kLocalVHeads + hv) * kChunk + col]; + } + s_a_aqk(row, col) = value; + } + for (int index = tid; index < kHeadDim * kChunk; index += kThreads) { + const int kk = index / kChunk; + const int row = index % kChunk; + bf16 value = bf16(0.0f); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + value = kg_bf16[(token * kLocalVHeads + hv) * kHeadDim + kk]; + } + s_a_kg(kk, row) = value; + } + const int last_token = seq * seq_len + chunk_start + valid_rows - 1; + const float chunk_decay = exp2f(g[last_token * kLocalVHeads + hv]); + __syncthreads(); + + if (warp == 0) { + completion.producer_acquire(producer_state); + cutlass::arch::fence_view_async_shared(); + kerutils::utcmma_ss( + mma_64x64, + s_a_aqk, + s_vnew, + t_acc_64x64_upper, + false); + kerutils::tcgen05_after_thread_sync(); + kerutils::utcmma_ss( + mma_128x64, + s_a_kg, + s_vnew, + t_acc_128x64_state, + true); + completion.producer_commit(producer_state); + ++producer_state; + } + completion.consumer_wait(consumer_state); + kerutils::tcgen05_after_thread_sync(); + + { + auto tiled_t2r = make_tmem_copy( + SM100_TMEM_LOAD_16dp256b8x{}, t_acc_64x64_upper); + auto thr_t2r = tiled_t2r.get_slice(tid); + auto t_src = thr_t2r.partition_S(t_acc_64x64_upper); + auto t_coord = thr_t2r.partition_D(t_c64); + auto r_acc = make_tensor(shape(t_coord)); + copy(tiled_t2r, t_src, r_acc); + cutlass::arch::fence_view_async_tmem_load(); + CUTE_UNROLL + for (int item = 0; item < size(r_acc); ++item) { + const auto coord = t_coord(item); + const int row = static_cast(get<0>(coord)); + const int vv = static_cast(get<1>(coord)); + if (row < valid_rows) { + const int token = seq * seq_len + chunk_start + row; + out[(token * kLocalVHeads + hv) * kHeadDim + v_base + vv] = + __float2bfloat16_rn(r_acc(item)); + } + } + } + + { + auto tiled_t2r = make_tmem_copy( + SM100_TMEM_LOAD_32dp32b16x{}, t_acc_128x64_state); + auto thr_t2r = tiled_t2r.get_slice(tid); + auto t_src = thr_t2r.partition_S(t_acc_128x64_state); + auto t_coord = thr_t2r.partition_D(t_c128); + auto r_acc = make_tensor(shape(t_coord)); + copy(tiled_t2r, t_src, r_acc); + cutlass::arch::fence_view_async_tmem_load(); + const bool is_last_chunk = chunk + 1 == chunk_count; + CUTE_UNROLL + for (int item = 0; item < size(r_acc); ++item) { + const auto coord = t_coord(item); + const int kk = static_cast(get<0>(coord)); + const int vv = static_cast(get<1>(coord)); + const float updated = + r_acc(item) + chunk_decay * static_cast(s_state(vv, kk)); + const bf16 quantized = bf16(updated); + s_state(vv, kk) = quantized; + if (is_last_chunk) { + final_state[state_global_base + kk * kHeadDim + v_base + vv] = + static_cast(quantized); + } + } + } + cutlass::arch::fence_view_async_shared(); + release_ss_mma_result(completion, consumer_state); + } + + __syncthreads(); + if (warp == 0) { + tmem_allocator.free(shared.tmem_base_ptr, kTmemColumns); + } +} + +template +inline void launch_qwen35_chunk_state_output_sm100_ss( + cudaStream_t stream, + const __nv_bfloat16* q_norm, + const float* g, + const __nv_bfloat16* Aqk, + const __nv_bfloat16* w, + const __nv_bfloat16* u, + const __nv_bfloat16* kg, + const float* initial_state, + __nv_bfloat16* out, + float* final_state, + int batch_size, + int seq_len, + int qk_heads, + bool has_initial_state) { + auto kernel_fn = &qwen35_chunk_state_output_sm100_ss_kernel; + constexpr size_t shared_bytes = + sizeof(Qwen35ChunkStateOutputSm100SsShared); + cudaFuncSetAttribute( + kernel_fn, + cudaFuncAttributeMaxDynamicSharedMemorySize, + shared_bytes); + const int grid = + batch_size * kLocalVHeads * (kHeadDim / kValueTile); + kernel_fn<<>>( + q_norm, + g, + Aqk, + w, + u, + kg, + initial_state, + out, + final_state, + batch_size, + seq_len, + qk_heads, + has_initial_state); +} + +} // namespace cula::qwen35::prefill::kernel::sm100_ss + +#endif // CULA_SM100_ENABLED From 51d6ab83af30547f03e81bafc4869d3528b6406f Mon Sep 17 00:00:00 2001 From: Xinhao Wei Date: Mon, 3 Aug 2026 16:09:47 +0000 Subject: [PATCH 35/35] perf(qwen35): specialize compact scalar gates on SM100 --- csrc/kda/sm100/fwd_helpers.hpp | 83 ++++----- csrc/kda/sm100/kda_fwd_common.cuh | 20 +- csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp | 48 ++++- .../sm100/kda_fwd_intra_mainloop_sm100.hpp | 87 +++++++-- .../sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp | 30 ++- .../kda_fwd_recomp_w_u_mainloop_sm100.hpp | 173 +++++++++++++----- csrc/kda/sm100/kda_fwd_sm100.cu | 10 + 7 files changed, 325 insertions(+), 126 deletions(-) diff --git a/csrc/kda/sm100/fwd_helpers.hpp b/csrc/kda/sm100/fwd_helpers.hpp index da9b5ae6..336387df 100644 --- a/csrc/kda/sm100/fwd_helpers.hpp +++ b/csrc/kda/sm100/fwd_helpers.hpp @@ -16,7 +16,10 @@ #include +#include + #include "kerutils/kerutils.cuh" +#include "kda/sm100/kda_fwd_common.cuh" namespace kda::sm100 { @@ -26,6 +29,21 @@ using ku::nvbf16x4; using ku::store_128b; using namespace cute; +template +CUTE_DEVICE void +gate_exp2_float4(float2& s1, float2& s2) { + if constexpr (std::is_same_v, ScalarGateView>) { + const float scale = exp2f(s1.x); + s1 = make_float2(scale, scale); + s2 = make_float2(scale, scale); + } else { + s1.x = exp2f(s1.x); + s1.y = exp2f(s1.y); + s2.x = exp2f(s2.x); + s2.y = exp2f(s2.y); + } +} + // ============================================================ // Forward Prologue: B-matrix (SMEM) helper functions // ============================================================ @@ -104,10 +122,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_0_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_0_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -117,10 +132,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_1_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_1_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -130,10 +142,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_2_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_2_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -143,10 +152,7 @@ fwd_setup_kg_col0_4out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -193,10 +199,7 @@ fwd_setup_kg_col1_3out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_1_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_1_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -206,10 +209,7 @@ fwd_setup_kg_col1_3out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_2_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_2_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -219,10 +219,7 @@ fwd_setup_kg_col1_3out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -266,10 +263,7 @@ fwd_setup_kg_col2_2out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_2_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_2_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -279,10 +273,7 @@ fwd_setup_kg_col2_2out( { float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], g_a); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], g_b); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, kf_a); reinterpret_cast(&res)[1] = float2_mul(s2, kf_b); @@ -317,10 +308,7 @@ fwd_setup_kg_col3_1out(G_TENSOR& sG, K_TENSOR& sK, KG_TENSOR& sKG_intra, int idx // intra(3,3): exp2(g_first_3 - g[x]) * K[x] float2 s1 = float2_sub(reinterpret_cast(&g_first_3_local)[0], reinterpret_cast(&g)[0]); float2 s2 = float2_sub(reinterpret_cast(&g_first_3_local)[1], reinterpret_cast(&g)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); float4 res; reinterpret_cast(&res)[0] = float2_mul(s1, __bfloat1622float2(k.a)); reinterpret_cast(&res)[1] = float2_mul(s2, __bfloat1622float2(k.b)); @@ -365,10 +353,7 @@ fwd_setup_A_inter_intra_all( float4 g_ref = *reinterpret_cast(&sG(g_first_row, y)); float2 s1 = float2_sub(g_a, reinterpret_cast(&g_ref)[0]); float2 s2 = float2_sub(g_b, reinterpret_cast(&g_ref)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); reinterpret_cast(&res_inter[i * 4])[0] = float2_mul(s1, va); reinterpret_cast(&res_inter[i * 4])[1] = float2_mul(s2, vb); } @@ -377,10 +362,7 @@ fwd_setup_A_inter_intra_all( float4 g_ref = *reinterpret_cast(&sG(g_half_row, y)); float2 s1 = float2_sub(g_a, reinterpret_cast(&g_ref)[0]); float2 s2 = float2_sub(g_b, reinterpret_cast(&g_ref)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); reinterpret_cast(&res_intra[i * 4])[0] = float2_mul(s1, va); reinterpret_cast(&res_intra[i * 4])[1] = float2_mul(s2, vb); } @@ -417,10 +399,7 @@ fwd_setup_A_inter_all( float4 g_ref = *reinterpret_cast(&sG(g_first_row, y)); float2 s1 = float2_sub(reinterpret_cast(&g)[0], reinterpret_cast(&g_ref)[0]); float2 s2 = float2_sub(reinterpret_cast(&g)[1], reinterpret_cast(&g_ref)[1]); - s1.x = exp2f(s1.x); - s1.y = exp2f(s1.y); - s2.x = exp2f(s2.x); - s2.y = exp2f(s2.y); + gate_exp2_float4(s1, s2); reinterpret_cast(&res_inter[i * 4])[0] = float2_mul(s1, va); reinterpret_cast(&res_inter[i * 4])[1] = float2_mul(s2, vb); } diff --git a/csrc/kda/sm100/kda_fwd_common.cuh b/csrc/kda/sm100/kda_fwd_common.cuh index 6bc227a0..ad5a8f46 100644 --- a/csrc/kda/sm100/kda_fwd_common.cuh +++ b/csrc/kda/sm100/kda_fwd_common.cuh @@ -18,14 +18,32 @@ namespace kda::sm100 { +// Presents a compact per-row scalar gate as the float4-addressable 2-D view +// used by the existing gated Q/K helpers. Each row stores four identical +// values; all logical K columns intentionally alias that four-float record. +struct ScalarGateView { + float* ptr; + + __device__ __forceinline__ float& + operator()(int row, int) const { + return ptr[row * 4]; + } +}; + // KDA forward kernels // KDA forward intra-chunk kernel void run_kda_fwd_intra_sm100(KDA_fwd_intra_params& params, cudaStream_t stream); +void +run_kda_fwd_intra_sm100_qwen_scalar_g(KDA_fwd_intra_params& params, cudaStream_t stream); + // KDA forward recompute W & U kernel void run_kda_fwd_recomp_w_u_sm100(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream); -} // namespace kda::sm100 \ No newline at end of file +void +run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream); + +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp index 3689bee3..2679772e 100644 --- a/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_kernel_sm100.hpp @@ -129,7 +129,9 @@ struct KdaChunkFwdIntraKernelSm100 { if (warp_idx == 0 && lane_predicate) { cute::prefetch_tma_descriptor(tma_params.tma_q.get_tma_descriptor()); cute::prefetch_tma_descriptor(tma_params.tma_k.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + if constexpr (!Mainloop::ScalarG) { + cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + } } // Allocate TMEM (warp 0 only) @@ -144,9 +146,10 @@ struct KdaChunkFwdIntraKernelSm100 { // === Unified TMA load pipeline: Q + K + G === typename PipelineQKG::Params qkg_load_pipe_params; - qkg_load_pipe_params.transaction_bytes = sizeof(ku::bf16) * cosize_v + // Q - sizeof(ku::bf16) * cosize_v + // K - sizeof(float) * cosize_v; // G + qkg_load_pipe_params.transaction_bytes = + sizeof(ku::bf16) * cosize_v + // Q + sizeof(ku::bf16) * cosize_v + // K + (Mainloop::ScalarG ? 0 : sizeof(float) * cosize_v); qkg_load_pipe_params.is_leader = lane_predicate && (role == WarpRole::Load); qkg_load_pipe_params.num_consumers = NumCudaCoreThreads; @@ -337,10 +340,16 @@ run_kda_fwd_intra_sm100_impl_dispatch(KDA_fwd_intra_params& params, cudaStream_t make_tensor(make_gmem_ptr((ku::bf16*)params.k_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); - auto tma_G = cute::make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), - typename Kernel::SmemLayoutInputFP32{}); + auto tma_G = [&]() { + if constexpr (Kernel::Mainloop::ScalarG) { + return 0; + } else { + return cute::make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), + typename Kernel::SmemLayoutInputFP32{}); + } + }(); // --- Pack TMA params --- typename Kernel:: @@ -374,11 +383,30 @@ run_kda_fwd_intra_sm100_impl(KDA_fwd_intra_params& params, cudaStream_t stream) BOOL_SWITCH(params.unified_gref, kUnifiedGRef, [&] { // Currently we hardcode RoundingTF32=false to align with FLA implementation, the precision is enough using Kernel = KdaChunkFwdIntraKernelSm100< - KdaChunkFwdIntraMainloopSm100>; + KdaChunkFwdIntraMainloopSm100< + kUseTF32Inverse, + /*RoundingTF32=*/false, + kUnifiedGRef, + /*ScalarG=*/false, + BetaType>>; run_kda_fwd_intra_sm100_impl_dispatch(params, stream); }); }); }); } -} // namespace kda::sm100 \ No newline at end of file +// Qwen GDN uses one scalar gate per token/value-head. Keep this entrypoint +// separate from the public vector-G dispatcher so the generic ABI and its +// template combinations remain unchanged. +inline void +run_kda_fwd_intra_sm100_qwen_scalar_g_impl(KDA_fwd_intra_params& params, cudaStream_t stream) { + using Kernel = KdaChunkFwdIntraKernelSm100>; + run_kda_fwd_intra_sm100_impl_dispatch(params, stream); +} + +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp index 55cdc686..1b6bb999 100644 --- a/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_intra_mainloop_sm100.hpp @@ -48,6 +48,7 @@ template < bool UseTF32Inverse_ = true, bool RoundingTF32_ = false, bool UnifiedGRef_ = false, + bool ScalarG_ = false, typename ElementBeta_ = float> struct KdaChunkFwdIntraMainloopSm100 { // ===================== Tile / Buffer Constants ===================== @@ -76,6 +77,7 @@ struct KdaChunkFwdIntraMainloopSm100 { // This makes inter and intra A-matrices identical, allowing the intra A-matrix to be skipped entirely. // Saves 50% of A-matrix exp2f computation and one TMEM store per k-iteration. static constexpr bool UnifiedGRef = UnifiedGRef_; + static constexpr bool ScalarG = ScalarG_; using ElementBeta = ElementBeta_; // double buffer in TMEM, overlap prologue A matrix with MMA @@ -194,12 +196,25 @@ struct KdaChunkFwdIntraMainloopSm100 { GmemLayoutAtom{}, Layout>>{})); // Val layout, 8 or 16 vals per store + struct VectorGateStorage { + array_aligned> g[StagesLoad]; + }; + + struct ScalarGateStorage { + // Four identical floats per row preserve the float4 load contract of + // the existing helpers while all logical K columns alias this record. + array_aligned g[StagesAcc]; + }; + + using GateStorage = std::conditional_t; + // ===================== Shared Memory Plan ===================== struct SharedMemoryPlan { - // Q, K, G double buffer + // Q/K use the TMA pipeline. Generic vector-G keeps a matching staged + // matrix; Qwen scalar-G keeps four duplicated floats per token row. array_aligned> q[StagesLoad]; // 12KB array_aligned> k[StagesLoad]; // 12KB - array_aligned> g[StagesLoad]; // 24KB + GateStorage gate; // Gated MMA K^T, double buffer struct { @@ -331,6 +346,13 @@ struct KdaChunkFwdIntraMainloopSm100 { int seq_len = cu_seqlens_ptr[batch_idx + 1] - cu_seqlens_ptr[batch_idx]; int sub_seq_len = min(TileT, seq_len - tile_idx * TileT); + // The Qwen specialization publishes compact scalar-g together + // with beta. Keep this stage until the tile epilogue so all four + // K slices can consume the expanded shared views. + if constexpr (ScalarG) { + beta_pipeline.consumer_wait(beta_pipe_state_read); + } + constexpr int kg_offset = SubTileT * TileK; // stride between sub_tile buffers CUTE_NO_UNROLL @@ -345,7 +367,16 @@ struct KdaChunkFwdIntraMainloopSm100 { // Step 2: Create SMEM tensor views for this buffer slot // ============================================================ Tensor sK = make_tensor(make_smem_ptr(shared_plan->k[buf_load_idx].data()), SmemLayoutInputBF16{}); - Tensor sG = make_tensor(make_smem_ptr(shared_plan->g[buf_load_idx].data()), SmemLayoutInputFP32{}); + auto sG = [&]() { + if constexpr (ScalarG) { + return ScalarGateView{ + shared_plan->gate.g[beta_pipe_state_read.index()].data()}; + } else { + return make_tensor( + make_smem_ptr(shared_plan->gate.g[buf_load_idx].data()), + SmemLayoutInputFP32{}); + } + }(); qkg_inter_pipeline.producer_acquire(qkg_inter_pipe_state_write); int buf_idx = qkg_inter_pipe_state_write.index(); @@ -484,7 +515,9 @@ struct KdaChunkFwdIntraMainloopSm100 { // and beta data before waiting for MMA, overlapping independent waits. kk_inv_pipeline.producer_acquire(kk_inv_pipe_state_write); - beta_pipeline.consumer_wait(beta_pipe_state_read); + if constexpr (!ScalarG) { + beta_pipeline.consumer_wait(beta_pipe_state_read); + } qk_done_pipeline.consumer_wait(qk_done_pipe_state_read); int buf_acc_idx = qk_done_pipe_state_read.index(); @@ -729,8 +762,15 @@ struct KdaChunkFwdIntraMainloopSm100 { make_coord(token_offset, _0{}, _0{}), tma_params.tma_q.get_tma_tensor(tma_params.shape_qk)); Tensor mK = domain_offset( make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_qk)); - Tensor mG = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + auto mG = [&]() { + if constexpr (ScalarG) { + return 0; + } else { + return domain_offset( + make_coord(token_offset, _0{}, _0{}), + tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + } + }(); // TMA load body (Q, K, G — unified pipeline, single barrier per stage) CUTE_NO_UNROLL @@ -738,20 +778,28 @@ struct KdaChunkFwdIntraMainloopSm100 { int buf_idx = qkg_load_pipe_state_write.index(); Tensor sQ = make_tensor(make_smem_ptr(shared_plan->q[buf_idx].data()), SmemLayoutInputBF16{}); Tensor sK = make_tensor(make_smem_ptr(shared_plan->k[buf_idx].data()), SmemLayoutInputBF16{}); - Tensor sG = make_tensor(make_smem_ptr(shared_plan->g[buf_idx].data()), SmemLayoutInputFP32{}); - // GVA: K and Q are sliced by qk_head_idx; G is sliced by head_idx (v-head). + // GVA: K and Q are sliced by qk_head_idx. Generic vector-G + // is sliced by v-head; Qwen scalar-G arrives via the aux + // pipeline and therefore has no TMA transfer here. Tensor gK = local_tile( mK(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); - Tensor gG = local_tile( - mG(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); Tensor gQ = local_tile( mQ(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, k_idx)); // Single acquire for all three TMA copies qkg_load_pipeline.producer_acquire(qkg_load_pipe_state_write); auto& barrier = *qkg_load_pipeline.producer_get_barrier(qkg_load_pipe_state_write); - ku::launch_tma_copy(tma_params.tma_g, gG, sG, barrier); + if constexpr (!ScalarG) { + Tensor sG = make_tensor( + make_smem_ptr(shared_plan->gate.g[buf_idx].data()), + SmemLayoutInputFP32{}); + Tensor gG = local_tile( + mG(_, _, head_idx), + make_shape(Int{}, Int{}), + make_coord(tile_idx, k_idx)); + ku::launch_tma_copy(tma_params.tma_g, gG, sG, barrier); + } ku::launch_tma_copy(tma_params.tma_k, gK, sK, barrier); ku::launch_tma_copy(tma_params.tma_q, gQ, sQ, barrier); ++qkg_load_pipe_state_write; @@ -911,11 +959,22 @@ struct KdaChunkFwdIntraMainloopSm100 { // Beta loading body beta_pipeline.producer_acquire(beta_pipe_state_write); if (thread_idx < TileT) { + const int token = token_offset + tile_idx * TileT + thread_idx; shared_plan->beta_smem[beta_pipe_state_write.index()][thread_idx] = (thread_idx < sub_seq_len) - ? float(reinterpret_cast( - params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h_v + head_idx]) + ? float(reinterpret_cast(params.beta_ptr)[token * params.h_v + head_idx]) : float(0); + if constexpr (ScalarG) { + const float gate = + (thread_idx < sub_seq_len) + ? reinterpret_cast(params.g_ptr)[token * params.h_v + head_idx] + : 0.0f; + float* gate4 = shared_plan->gate.g[beta_pipe_state_write.index()].data() + thread_idx * 4; + gate4[0] = gate; + gate4[1] = gate; + gate4[2] = gate; + gate4[3] = gate; + } } fence_view_async_shared(); beta_pipeline.producer_commit(beta_pipe_state_write); @@ -924,4 +983,4 @@ struct KdaChunkFwdIntraMainloopSm100 { } }; -} // namespace kda::sm100 \ No newline at end of file +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp index 2cae6c04..482290dd 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_kernel_sm100.hpp @@ -152,7 +152,9 @@ struct KdaChunkFwdRecompWUKernelSm100 { cute::prefetch_tma_descriptor(tma_params.tma_akk.get_tma_descriptor()); cute::prefetch_tma_descriptor(tma_params.tma_k.get_tma_descriptor()); cute::prefetch_tma_descriptor(tma_params.tma_v.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + if constexpr (!Mainloop::ScalarG) { + cute::prefetch_tma_descriptor(tma_params.tma_g.get_tma_descriptor()); + } if constexpr (StoreQG) { cute::prefetch_tma_descriptor(tma_params.tma_q.get_tma_descriptor()); } @@ -453,10 +455,16 @@ run_kda_fwd_recomp_w_u_sm100_impl_dispatch(KDA_fwd_recomp_w_u_params& params, cu make_tensor(make_gmem_ptr((bf16*)params.k_ptr), make_layout(shape_QK, stride_QK)), typename Kernel::SmemLayoutInputBF16{}); - auto tma_G = cute::make_tma_copy( - SM90_TMA_LOAD{}, - make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), - typename Kernel::SmemLayoutInputFP32{}); + auto tma_G = [&]() { + if constexpr (Kernel::Mainloop::ScalarG) { + return 0; + } else { + return cute::make_tma_copy( + SM90_TMA_LOAD{}, + make_tensor(make_gmem_ptr((float*)params.g_ptr), make_layout(shape_VG, stride_VG)), + typename Kernel::SmemLayoutInputFP32{}); + } + }(); auto tma_Akk = cute::make_tma_copy( SM90_TMA_LOAD{}, @@ -502,10 +510,18 @@ inline void run_kda_fwd_recomp_w_u_sm100_impl(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { BETA_TYPE_SWITCH(params.is_beta_bf16, BetaType, [&] { BOOL_SWITCH(params.store_qg, kStoreQG, [&] { - using Kernel = KdaChunkFwdRecompWUKernelSm100>; + using Kernel = KdaChunkFwdRecompWUKernelSm100< + KdaChunkFwdRecompWUMainloopSm100>; run_kda_fwd_recomp_w_u_sm100_impl_dispatch(params, stream); }); }); } -} // namespace kda::sm100 \ No newline at end of file +inline void +run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g_impl(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { + using Kernel = KdaChunkFwdRecompWUKernelSm100< + KdaChunkFwdRecompWUMainloopSm100>; + run_kda_fwd_recomp_w_u_sm100_impl_dispatch(params, stream); +} + +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp index d702e46f..30568877 100644 --- a/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp +++ b/csrc/kda/sm100/kda_fwd_recomp_w_u_mainloop_sm100.hpp @@ -35,7 +35,7 @@ struct KdaChunkFwdRecompWUSm100NamedBarriers { // constants, and the persistent loop bodies for each warp role. // The Kernel struct is templated on this Mainloop. // =================================================================== -template +template struct KdaChunkFwdRecompWUMainloopSm100 { // ===================== Tile / Buffer Constants ===================== static constexpr int TileT = 64; @@ -49,6 +49,7 @@ struct KdaChunkFwdRecompWUMainloopSm100 { static constexpr int StagesQ = 1; static constexpr bool StoreQG = StoreQG_; + static constexpr bool ScalarG = ScalarG_; using ElementBeta = ElementBeta_; // TODO: try optimization with tcgen05.mma.ws @@ -160,13 +161,22 @@ struct KdaChunkFwdRecompWUMainloopSm100 { struct QSmemBufferDisabled {}; // empty, zero-cost using QSmemBuffer = cute::conditional_t; + struct VectorGateStorage { + array_aligned> g[StagesLoadStore]; + }; + struct ScalarGateStorage { + array_aligned g[StagesA]; + }; + using GateStorage = cute::conditional_t; + struct SharedMemoryPlan { // Akk, single buffer array_aligned> akk[StagesA]; // 16KB - // K, V, G double buffer + // K/V are double buffered. Generic vector-G keeps the full staged + // matrix; Qwen scalar-G keeps four duplicated floats per token row. array_aligned> k[StagesLoadStore]; // 32KB array_aligned> v[StagesLoadStore]; // 32KB - array_aligned> g[StagesLoadStore]; // 64KB + GateStorage gate; // Q double buffer (only present when StoreQG=true) QSmemBuffer q_buf; // MMA B-operand staging: K_proc/V_proc after prologue, [N=TileK, K=TileT] MN-major, double buffer @@ -326,9 +336,19 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } - g_pipeline.consumer_wait(g_pipe_state_read); - Tensor sG = - make_tensor(make_smem_ptr(shared_plan->g[g_pipe_state_read.index()].data()), SmemLayoutInputFP32{}); + if constexpr (!ScalarG) { + g_pipeline.consumer_wait(g_pipe_state_read); + } + auto sG = [&]() { + if constexpr (ScalarG) { + return ScalarGateView{ + shared_plan->gate.g[beta_pipe_state_read.index()].data()}; + } else { + return make_tensor( + make_smem_ptr(shared_plan->gate.g[g_pipe_state_read.index()].data()), + SmemLayoutInputFP32{}); + } + }(); // Load G with same 16x64 column mapping as K (two float4 per iteration) #pragma unroll for (int ti = 0; ti < TileT / 16; ++ti) { @@ -368,15 +388,27 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // lo half (cols y..y+3): k_reg a01, a23 + g_reg lo float2 kf_01 = __bfloat1622float2(k_reg[ti][k_yi].a01); float2 kf_23 = __bfloat1622float2(k_reg[ti][k_yi].a23); - float2 g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; - float2 g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + float2 g_01; + float2 g_23; + float2 g_45; + float2 g_67; + if constexpr (ScalarG) { + const float scale = exp2f(g_reg[ti][k_yi][0].x); + g_01 = make_float2(scale, scale); + g_23 = make_float2(scale, scale); + g_45 = make_float2(scale, scale); + g_67 = make_float2(scale, scale); + } else { + g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; + g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; + g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; + } float2 res_01 = float2_mul(float2_mul(kf_01, beta2), g_01); float2 res_23 = float2_mul(float2_mul(kf_23, beta2), g_23); // hi half (cols y+4..y+7): k_reg a45, a67 + g_reg hi float2 kf_45 = __bfloat1622float2(k_reg[ti][k_yi].a45); float2 kf_67 = __bfloat1622float2(k_reg[ti][k_yi].a67); - float2 g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; - float2 g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; float2 res_45 = float2_mul(float2_mul(kf_45, beta2), g_45); float2 res_67 = float2_mul(float2_mul(kf_67, beta2), g_67); // Single 128-bit store @@ -431,23 +463,36 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // lo half (cols y..y+3): k_reg a01, a23 float2 kf_01 = __bfloat1622float2(k_reg[ti][k_yi].a01); float2 kf_23 = __bfloat1622float2(k_reg[ti][k_yi].a23); - float2 gd_01 = { - exp2f(g_last_reg[k_yi][0].x - g_reg[ti][k_yi][0].x), - exp2f(g_last_reg[k_yi][0].y - g_reg[ti][k_yi][0].y)}; - float2 gd_23 = { - exp2f(g_last_reg[k_yi][0].z - g_reg[ti][k_yi][0].z), - exp2f(g_last_reg[k_yi][0].w - g_reg[ti][k_yi][0].w)}; + float2 gd_01; + float2 gd_23; + float2 gd_45; + float2 gd_67; + if constexpr (ScalarG) { + const float scale = exp2f( + g_last_reg[k_yi][0].x - g_reg[ti][k_yi][0].x); + gd_01 = make_float2(scale, scale); + gd_23 = make_float2(scale, scale); + gd_45 = make_float2(scale, scale); + gd_67 = make_float2(scale, scale); + } else { + gd_01 = { + exp2f(g_last_reg[k_yi][0].x - g_reg[ti][k_yi][0].x), + exp2f(g_last_reg[k_yi][0].y - g_reg[ti][k_yi][0].y)}; + gd_23 = { + exp2f(g_last_reg[k_yi][0].z - g_reg[ti][k_yi][0].z), + exp2f(g_last_reg[k_yi][0].w - g_reg[ti][k_yi][0].w)}; + gd_45 = { + exp2f(g_last_reg[k_yi][1].x - g_reg[ti][k_yi][1].x), + exp2f(g_last_reg[k_yi][1].y - g_reg[ti][k_yi][1].y)}; + gd_67 = { + exp2f(g_last_reg[k_yi][1].z - g_reg[ti][k_yi][1].z), + exp2f(g_last_reg[k_yi][1].w - g_reg[ti][k_yi][1].w)}; + } float2 res_01 = float2_mul(kf_01, gd_01); float2 res_23 = float2_mul(kf_23, gd_23); // hi half (cols y+4..y+7): k_reg a45, a67 float2 kf_45 = __bfloat1622float2(k_reg[ti][k_yi].a45); float2 kf_67 = __bfloat1622float2(k_reg[ti][k_yi].a67); - float2 gd_45 = { - exp2f(g_last_reg[k_yi][1].x - g_reg[ti][k_yi][1].x), - exp2f(g_last_reg[k_yi][1].y - g_reg[ti][k_yi][1].y)}; - float2 gd_67 = { - exp2f(g_last_reg[k_yi][1].z - g_reg[ti][k_yi][1].z), - exp2f(g_last_reg[k_yi][1].w - g_reg[ti][k_yi][1].w)}; float2 res_45 = float2_mul(kf_45, gd_45); float2 res_67 = float2_mul(kf_67, gd_67); // Single 128-bit store @@ -468,8 +513,10 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } } - g_pipeline.consumer_release(g_pipe_state_read); - ++g_pipe_state_read; + if constexpr (!ScalarG) { + g_pipeline.consumer_release(g_pipe_state_read); + ++g_pipe_state_read; + } // Ensure all 128 prologue threads have finished writing sKG_out cutlass::arch::NamedBarrier::arrive_and_wait( @@ -534,15 +581,27 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // lo half (cols y..y+3) float2 qf_01 = __bfloat1622float2(q_reg[ti][k_yi].a01); float2 qf_23 = __bfloat1622float2(q_reg[ti][k_yi].a23); - float2 g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; - float2 g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + float2 g_01; + float2 g_23; + float2 g_45; + float2 g_67; + if constexpr (ScalarG) { + const float scale = exp2f(g_reg[ti][k_yi][0].x); + g_01 = make_float2(scale, scale); + g_23 = make_float2(scale, scale); + g_45 = make_float2(scale, scale); + g_67 = make_float2(scale, scale); + } else { + g_01 = {exp2f(g_reg[ti][k_yi][0].x), exp2f(g_reg[ti][k_yi][0].y)}; + g_23 = {exp2f(g_reg[ti][k_yi][0].z), exp2f(g_reg[ti][k_yi][0].w)}; + g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; + g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; + } float2 res_01 = float2_mul(qf_01, g_01); float2 res_23 = float2_mul(qf_23, g_23); // hi half (cols y+4..y+7) float2 qf_45 = __bfloat1622float2(q_reg[ti][k_yi].a45); float2 qf_67 = __bfloat1622float2(q_reg[ti][k_yi].a67); - float2 g_45 = {exp2f(g_reg[ti][k_yi][1].x), exp2f(g_reg[ti][k_yi][1].y)}; - float2 g_67 = {exp2f(g_reg[ti][k_yi][1].z), exp2f(g_reg[ti][k_yi][1].w)}; float2 res_45 = float2_mul(qf_45, g_45); float2 res_67 = float2_mul(qf_67, g_67); // Single 128-bit store @@ -919,8 +978,15 @@ struct KdaChunkFwdRecompWUMainloopSm100 { make_coord(token_offset, _0{}, _0{}), tma_params.tma_k.get_tma_tensor(tma_params.shape_qk)); Tensor mV = domain_offset( make_coord(token_offset, _0{}, _0{}), tma_params.tma_v.get_tma_tensor(tma_params.shape_vg)); - Tensor mG = domain_offset( - make_coord(token_offset, _0{}, _0{}), tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + auto mG = [&]() { + if constexpr (ScalarG) { + return 0; + } else { + return domain_offset( + make_coord(token_offset, _0{}, _0{}), + tma_params.tma_g.get_tma_tensor(tma_params.shape_vg)); + } + }(); Tensor mA = domain_offset( make_coord(token_offset, _0{}, _0{}), tma_params.tma_akk.get_tma_tensor(tma_params.shape_Akk)); @@ -958,26 +1024,37 @@ struct KdaChunkFwdRecompWUMainloopSm100 { make_smem_ptr(shared_plan->k[k_pipe_state_write.index()].data()), SmemLayoutInputBF16{}); Tensor sV = make_tensor( make_smem_ptr(shared_plan->v[v_pipe_state_write.index()].data()), SmemLayoutInputBF16{}); - Tensor sG = make_tensor( - make_smem_ptr(shared_plan->g[g_pipe_state_write.index()].data()), SmemLayoutInputFP32{}); - // GVA slicing: K uses qk_head_idx; V and G use the v-head index. + // GVA slicing: K uses qk_head_idx and V uses v-head. + // Generic vector-G is loaded by TMA; Qwen scalar-G is + // published by the aux pipeline instead. Tensor gK = local_tile( mK(_, _, qk_head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); Tensor gV = local_tile( mV(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); - Tensor gG = local_tile( - mG(_, _, head_idx), make_shape(Int{}, Int{}), make_coord(tile_idx, i_k)); // K: Load → Compute k_pipeline.producer_acquire(k_pipe_state_write); ku::launch_tma_copy(tma_params.tma_k, gK, sK, *k_pipeline.producer_get_barrier(k_pipe_state_write)); ++k_pipe_state_write; - // G: Load → Compute - g_pipeline.producer_acquire(g_pipe_state_write); - ku::launch_tma_copy(tma_params.tma_g, gG, sG, *g_pipeline.producer_get_barrier(g_pipe_state_write)); - ++g_pipe_state_write; + // G: Load → Compute (generic vector-G only) + if constexpr (!ScalarG) { + Tensor sG = make_tensor( + make_smem_ptr(shared_plan->gate.g[g_pipe_state_write.index()].data()), + SmemLayoutInputFP32{}); + Tensor gG = local_tile( + mG(_, _, head_idx), + make_shape(Int{}, Int{}), + make_coord(tile_idx, i_k)); + g_pipeline.producer_acquire(g_pipe_state_write); + ku::launch_tma_copy( + tma_params.tma_g, + gG, + sG, + *g_pipeline.producer_get_barrier(g_pipe_state_write)); + ++g_pipe_state_write; + } // V: Load → Compute v_pipeline.producer_acquire(v_pipe_state_write); @@ -1039,12 +1116,24 @@ struct KdaChunkFwdRecompWUMainloopSm100 { // ============================================================ beta_pipeline.producer_acquire(beta_pipe_state_write); if (thread_idx < TileT) { + const int token = token_offset + tile_idx * TileT + thread_idx; float beta_val = (thread_idx < sub_seq_len) - ? float(reinterpret_cast( - params.beta_ptr)[(token_offset + tile_idx * TileT + thread_idx) * params.h_v + head_idx]) + ? float(reinterpret_cast(params.beta_ptr)[token * params.h_v + head_idx]) : float(0); shared_plan->beta_smem[beta_pipe_state_write.index()][thread_idx] = beta_val; + if constexpr (ScalarG) { + const float gate = + (thread_idx < sub_seq_len) + ? reinterpret_cast(params.g_ptr)[token * params.h_v + head_idx] + : 0.0f; + float* gate4 = + shared_plan->gate.g[beta_pipe_state_write.index()].data() + thread_idx * 4; + gate4[0] = gate; + gate4[1] = gate; + gate4[2] = gate; + gate4[3] = gate; + } } fence_view_async_shared(); beta_pipeline.producer_commit(beta_pipe_state_write); @@ -1053,4 +1142,4 @@ struct KdaChunkFwdRecompWUMainloopSm100 { } }; -} // namespace kda::sm100 \ No newline at end of file +} // namespace kda::sm100 diff --git a/csrc/kda/sm100/kda_fwd_sm100.cu b/csrc/kda/sm100/kda_fwd_sm100.cu index edaaf0d9..1c10d7f2 100644 --- a/csrc/kda/sm100/kda_fwd_sm100.cu +++ b/csrc/kda/sm100/kda_fwd_sm100.cu @@ -23,9 +23,19 @@ run_kda_fwd_intra_sm100(KDA_fwd_intra_params& params, cudaStream_t stream) { kda::sm100::run_kda_fwd_intra_sm100_impl(params, stream); } +void +run_kda_fwd_intra_sm100_qwen_scalar_g(KDA_fwd_intra_params& params, cudaStream_t stream) { + kda::sm100::run_kda_fwd_intra_sm100_qwen_scalar_g_impl(params, stream); +} + void run_kda_fwd_recomp_w_u_sm100(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { kda::sm100::run_kda_fwd_recomp_w_u_sm100_impl(params, stream); } +void +run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g(KDA_fwd_recomp_w_u_params& params, cudaStream_t stream) { + kda::sm100::run_kda_fwd_recomp_w_u_sm100_qwen_scalar_g_impl(params, stream); +} + } // namespace kda::sm100