diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index e8ec53246..69f4a0cac 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -85,7 +85,6 @@ set(WORKSPACE_KERNEL_SRCS ${PROJECT_OP_SRC_BASE}/build_tree/op_kernel/build_tree_kernel.cpp ${PROJECT_OP_SRC_BASE}/lightning_indexer/op_kernel/lightning_indexer_kernel.cpp ${PROJECT_OP_SRC_BASE}/causal_conv1d_update/op_kernel/causal_conv1d_update.cpp - ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_kernel/causal_conv1d.cpp ${PROJECT_OP_SRC_BASE}/lora/op_kernel/sgemmc_expand_kernel.cpp ${PROJECT_OP_SRC_BASE}/lora/op_kernel/sgemmc_shrink_kernel.cpp ) @@ -132,11 +131,24 @@ ascendc_compile_definitions(mega_chunk_gdn_kernel PRIVATE -DGDN_C=128 ) +# PTO-ISA causal_conv1d kernel (depthwise conv + bias + SiLU), built like mega_chunk_gdn. +ascendc_library(causal_conv1d_kernel STATIC + ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_kernel/causal_conv1d.cpp +) +ascendc_include_directories(causal_conv1d_kernel PRIVATE + ${PROJECT_OP_SRC_BASE}/causal_conv1d/op_kernel + ${PROJECT_SOURCE_DIR}/third_party/pto-isa/include + ${ASCEND_INCLUDE_DIR} + ${ASCEND_INCLUDE_DIR}/experiment/runtime + ${ASCEND_INCLUDE_DIR}/experiment/msprof +) + # create shared library libsgl_kernel_npu.so add_library(${OP_PLUGIN_NAME} SHARED ${OP_SRCS}) target_link_libraries(${OP_PLUGIN_NAME} PRIVATE ${MEGA_CHUNK_GDN_KERNEL_TARGETS} + causal_conv1d_kernel workspace_kernel no_workspace_kernel torch_npu diff --git a/csrc/causal_conv1d/op_host/causal_conv1d.cpp b/csrc/causal_conv1d/op_host/causal_conv1d.cpp index cf76fbd24..b41d61548 100644 --- a/csrc/causal_conv1d/op_host/causal_conv1d.cpp +++ b/csrc/causal_conv1d/op_host/causal_conv1d.cpp @@ -1,405 +1,270 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See - * LICENSE in the root of the software repository for the full text of the License. - */ - -/*! - * \file causal_conv1d.cpp - * \brief causal_conv1d host-side implementation - */ +#include "causal_conv1d.h" -#include +#include #include -#include +#include +#include + #include "acl/acl.h" -#include -#include -#include -#include "torch_npu/csrc/core/npu/NPUStream.h" -#include "torch_npu/csrc/core/npu/DeviceUtils.h" #include "tiling/platform/platform_ascendc.h" -#include "stub/aclrtlaunch_causal_conv1d.h" -#include "defines.h" #include "torch_helper.h" -#include "common.h" -#include "causal_conv1d.h" -#include "../op_kernel/causal_conv1d_tiling_data.h" + +// Ring sizes the kernel is compiled for -- must match FOR_EACH_RING_SIZE in +// op_kernel/causal_conv1d.cpp. Each row is (ringSize, maxTileWidth); this one list +// drives the launch-stub declarations, the per-ring tile-width lookup, and the dispatch. +#define FOR_EACH_RING_SIZE(DO) DO(2, 4096) DO(4, 3072) DO(8, 1536) DO(16, 896) DO(32, 384) DO(64, 128) + +// AscendC emits one host launch stub (aclrtlaunch_) per kernel entry. A macro +// can't generate #include directives, so rather than pull in 24 generated headers we +// forward-declare the stubs from the ring-size list (their definitions come from the +// linked causal_conv1d_kernel lib; same approach as causal_conv1d_update). + +// EXEC_KERNEL_CMD launches via ACLRT_LAUNCH_KERNEL(name) -> the aclrtlaunch_ symbol. +#ifndef ACLRT_LAUNCH_KERNEL +#define ACLRT_LAUNCH_KERNEL(kernel_func) aclrtlaunch_##kernel_func +#endif +// clang-format off +// Launch-stub arg types = (blockDim, stream, then the kernel params with GM_ADDR -> +// void*); must mirror CONV_PARAMS / WB_PARAMS in op_kernel/causal_conv1d.cpp. +#define CONV_STUB_PARAMS \ + uint32_t, aclrtStream, void *, void *, void *, void *, void *, void *, void *, void *, uint32_t, uint32_t, \ + uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, int32_t +#define WRITEBACK_STUB_PARAMS \ + uint32_t, aclrtStream, void *, void *, void *, void *, void *, uint32_t, uint32_t, uint32_t, uint32_t, \ + uint32_t, uint32_t, uint32_t, uint32_t, int32_t +// Worker macro: one ring size -> the conv + writeback stub, each for half and bf16. +#define DECLARE_LAUNCH_STUBS(ringSize, maxTileWidth) \ + extern "C" uint32_t aclrtlaunch_causal_conv1d_rs##ringSize##_half(CONV_STUB_PARAMS); \ + extern "C" uint32_t aclrtlaunch_causal_conv1d_rs##ringSize##_bf16(CONV_STUB_PARAMS); \ + extern "C" uint32_t aclrtlaunch_causal_conv1d_wb_rs##ringSize##_half(WRITEBACK_STUB_PARAMS); \ + extern "C" uint32_t aclrtlaunch_causal_conv1d_wb_rs##ringSize##_bf16(WRITEBACK_STUB_PARAMS); +FOR_EACH_RING_SIZE(DECLARE_LAUNCH_STUBS) // expand the list -> all 24 launch-stub declarations +#undef DECLARE_LAUNCH_STUBS +#undef WRITEBACK_STUB_PARAMS +#undef CONV_STUB_PARAMS +// clang-format on namespace sglang { namespace npu_kernel { +namespace { // file-local helpers (internal linkage; avoids clashes with other ops) -constexpr uint32_t PADDING_BYTE = 32U; -constexpr int64_t DIM_ALIGN = 16; -constexpr int64_t MAX_DIM_TILE = 4096; -constexpr int32_t MAX_WIDTH = 4; -constexpr int32_t MIN_WIDTH = 2; -constexpr int64_t ASCENDC_RESERVED_WORKSPACE = 16 * 1024 * 1024; -constexpr uint32_t MAX_CAPTURE_NUM = 1024; - -constexpr uint32_t CAUSAL_CONV1D_TPL_RUN_MODE_FN = 0; -constexpr uint32_t CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE = 1; -constexpr uint32_t CAUSAL_CONV1D_TPL_WIDTH_2 = 1; -constexpr uint32_t CAUSAL_CONV1D_TPL_WIDTH_3 = 2; -constexpr uint32_t CAUSAL_CONV1D_TPL_WIDTH_4 = 3; -constexpr uint32_t CAUSAL_CONV1D_TPL_FN_PLAN_INVALID = 0; -constexpr uint32_t CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS = 1; -constexpr uint32_t CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD = 2; - -static uint32_t g_causalConv1dCaptureNum = 0; -static std::unordered_map g_causalConv1dCaptureMap; - -struct CausalConv1dTilingKey { - int64_t dim; - int64_t cuSeqlen; - int64_t seqLen; - int64_t batch; - int64_t inputMode; - int64_t width; - int64_t stateLen; - int64_t numCacheLines; - int64_t activationMode; - int64_t padSlotId; - int64_t runMode; - int64_t hasBias; - int64_t hasCacheIndices; - int64_t hasInitialState; - int64_t hasNumAccept; -}; +// Accumulator-ring size = smallest power of two >= width. The kernel templates on the +// ring size (compile-time) and takes the width at runtime; the host computes the ring +// size here and launches the matching variant, so any width <= ring size reuses it. +constexpr uint32_t roundUpToPow2(uint32_t width) +{ + uint32_t n = (width != 0u) ? width - 1u : 0u; + n |= n >> 1u; + n |= n >> 2u; + n |= n >> 4u; + n |= n >> 8u; + n |= n >> 16u; + return n + 1u; +} -struct CausalConv1dTilingKeyHash { - static inline std::size_t HashCombine(std::size_t seed, std::size_t value) - { - seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); - return seed; +// Per-ring channel-tile width -- MUST match the (ringSize, maxTileWidth) the kernel is +// compiled with. Architectures with larger UB, can use wider tiles. +// The host has no compile-time device-arch macro, so the caller picks the table at runtime +#define FOR_EACH_RING_SIZE_A5(DO) DO(2, 5120) DO(4, 4096) DO(8, 2048) DO(16, 1152) DO(32, 512) DO(64, 128) +#define MAX_WIDTH_CASE(ringSize, maxTileWidth) \ + case ringSize: \ + return maxTileWidth##u; +uint32_t maxTileWidthForRing(uint32_t ringSize, bool wideUb) +{ + if (wideUb) { + switch (ringSize) { + FOR_EACH_RING_SIZE_A5(MAX_WIDTH_CASE) + default: + return 0u; + } } - - std::size_t operator()(const CausalConv1dTilingKey &k) const - { - std::size_t h = 0; - h = HashCombine(h, static_cast(k.dim)); - h = HashCombine(h, static_cast(k.cuSeqlen)); - h = HashCombine(h, static_cast(k.seqLen)); - h = HashCombine(h, static_cast(k.batch)); - h = HashCombine(h, static_cast(k.inputMode)); - h = HashCombine(h, static_cast(k.width)); - h = HashCombine(h, static_cast(k.stateLen)); - h = HashCombine(h, static_cast(k.numCacheLines)); - h = HashCombine(h, static_cast(k.activationMode)); - h = HashCombine(h, static_cast(k.padSlotId)); - h = HashCombine(h, static_cast(k.runMode)); - h = HashCombine(h, static_cast(k.hasBias)); - h = HashCombine(h, static_cast(k.hasCacheIndices)); - h = HashCombine(h, static_cast(k.hasInitialState)); - h = HashCombine(h, static_cast(k.hasNumAccept)); - return h; + switch (ringSize) { + FOR_EACH_RING_SIZE(MAX_WIDTH_CASE) + default: + return 0u; } -}; - -namespace { +} +#undef MAX_WIDTH_CASE +#undef FOR_EACH_RING_SIZE_A5 -inline int64_t CeilDiv(int64_t x, int64_t y) +// Supported filter widths: any width in [2, 64], routed to the roundUpToPow2(width) +// ring variant. width > 64 would need ring 128, which does not fit UB. +constexpr bool isSupportedWidth(uint32_t width) { - return (x + y - 1) / y; + return width >= 2u && width <= 64u; } -inline int64_t AlignUp(int64_t value, int64_t align) +template +constexpr T ceil_div(T a, T b) { - return CeilDiv(value, align) * align; + return (a + b - 1) / b; } - -struct UpdateDimTileChoice { - int64_t baseDim = 0; - int64_t baseDimCnt = 0; - int64_t gridSize = 0; -}; - -// Mirror of the GE update-mode tiling policy (ChooseCanonicalUpdateBaseDimChoice): -// pick a baseDim from a fixed candidate set that ideally divides dim and makes the -// batch*baseDimCnt grid land as close as possible to (>=) the AIV core count, so the -// batch-parallel update kernel keeps all cores busy. Falls back to ceil-division when -// no candidate divides dim exactly. -UpdateDimTileChoice ChooseUpdateBaseDimChoice(int64_t batch, int64_t dim, int32_t numCores) +constexpr uint32_t round_up_128(uint32_t x) { - const int64_t candidates[] = {4096, 2048, 1024, 512, 384, 192}; - const int64_t coreNum = (numCores > 0) ? static_cast(numCores) : 1; - - auto chooseOnce = [&](bool requireExactDiv) -> UpdateDimTileChoice { - UpdateDimTileChoice bestOver; - int64_t bestOverGap = std::numeric_limits::max(); - UpdateDimTileChoice bestUnder; - - for (int64_t candBaseDim : candidates) { - if (candBaseDim <= 0) { - continue; - } - if (requireExactDiv && (dim % candBaseDim != 0)) { - continue; - } - const int64_t baseDimCnt = requireExactDiv ? (dim / candBaseDim) : CeilDiv(dim, candBaseDim); - const int64_t gridSize = batch * baseDimCnt; - if (gridSize <= 0) { - continue; - } - if (gridSize >= coreNum) { - const int64_t gap = gridSize - coreNum; - if (gap < bestOverGap) { - bestOver = {candBaseDim, baseDimCnt, gridSize}; - bestOverGap = gap; - } - } else if (gridSize > bestUnder.gridSize || - (gridSize == bestUnder.gridSize && candBaseDim < bestUnder.baseDim)) { - bestUnder = {candBaseDim, baseDimCnt, gridSize}; - } - } - return (bestOver.baseDim != 0) ? bestOver : bestUnder; - }; - - UpdateDimTileChoice result = chooseOnce(true); - if (result.baseDim == 0) { - result = chooseOnce(false); - } - return result; + return ((x + 127u) / 128u) * 128u; } -void ComputeTilingData(int64_t dim, int64_t cuSeqlen, int64_t seqLen, int64_t batch, int64_t inputMode, int64_t width, - int64_t stateLen, int64_t numCacheLines, int64_t activationMode, int64_t padSlotId, bool hasBias, - bool hasCacheIndices, bool hasInitialState, bool hasNumAccept, bool isBf16, int32_t numCores, - int64_t runMode, CausalConv1dTilingData &td) +std::pair tiling_causal_conv1d(uint64_t numCores, uint64_t batch, const uint64_t dim, + const uint64_t seqLength, const uint64_t width, + const uint64_t maxChannels) { - (void)padSlotId; - std::memset(&td, 0, sizeof(td)); - - td.dim = dim; - td.cuSeqlen = cuSeqlen; - td.seqLen = seqLen; - td.inputMode = inputMode; - td.width = width; - td.stateLen = stateLen; - td.numCacheLines = numCacheLines; - td.batch = batch; - td.activationMode = activationMode; - td.padSlotId = padSlotId; - td.hasBias = hasBias ? 1 : 0; - td.hasCacheIndices = hasCacheIndices ? 1 : 0; - td.hasInitialStateMode = hasInitialState ? 1 : 0; - td.hasInitStateWorkspace = hasInitialState ? 1 : 0; - td.hasNumAcceptedTokens = hasNumAccept ? 1 : 0; - - td.dtypeKey = isBf16 ? 0 : 1; - td.runModeKey = static_cast(runMode); - td.widthKey = (width == 2) ? CAUSAL_CONV1D_TPL_WIDTH_2 - : (width == 3) ? CAUSAL_CONV1D_TPL_WIDTH_3 - : CAUSAL_CONV1D_TPL_WIDTH_4; - - if (runMode == CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE) { - UpdateDimTileChoice choice = ChooseUpdateBaseDimChoice(batch, dim, numCores); - if (choice.baseDim <= 0 || choice.baseDimCnt <= 0) { - choice.baseDim = (dim > 0 && dim <= MAX_DIM_TILE) ? dim : MAX_DIM_TILE; - choice.baseDimCnt = (choice.baseDim > 0) ? CeilDiv(dim, choice.baseDim) : 1; - if (choice.baseDimCnt <= 0) { - choice.baseDimCnt = 1; - } - } - td.baseDim = choice.baseDim; - td.baseDimCnt = choice.baseDimCnt; - td.fnPlanKey = CAUSAL_CONV1D_TPL_FN_PLAN_INVALID; - td.tokenBlockSize = 0; - td.tokenBlockCnt = 0; - } else if (dim <= MAX_DIM_TILE && numCores > 0) { - td.baseDim = dim; - td.baseDimCnt = 1; - td.fnPlanKey = CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS; - int64_t tokenCoreBudget = numCores; - int64_t idealBlockSize = CeilDiv(cuSeqlen, tokenCoreBudget); - if (idealBlockSize <= 0) { - idealBlockSize = 1; - } - td.tokenBlockSize = idealBlockSize; - td.tokenBlockCnt = CeilDiv(cuSeqlen, td.tokenBlockSize); - } else { - td.baseDim = MAX_DIM_TILE; - td.baseDimCnt = CeilDiv(dim, td.baseDim); - td.fnPlanKey = CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD; - int64_t tokenCoreBudget = (numCores > 0) ? (numCores / td.baseDimCnt) : 1; - if (tokenCoreBudget <= 0) { - tokenCoreBudget = 1; + // 128 = Number of Vector Lanes in FP16/BF16 + constexpr uint64_t minChannels = 128; // Must divide maxChannels + + uint64_t gcdCoreBatch = std::gcd(numCores, batch); + numCores /= gcdCoreBatch; + batch /= gcdCoreBatch; + + uint64_t numChannels = ceil_div(dim, maxChannels); + uint64_t channelsPerTile = ceil_div(ceil_div(dim, numChannels), minChannels) * minChannels; + + uint64_t seqChunks = 1u; + double bestScore = std::numeric_limits::infinity(); + + numChannels = ceil_div(dim, channelsPerTile); + const uint64_t depthNumerator = batch * numChannels; + + const uint64_t uppBnd = numCores / std::gcd(numCores, numChannels); + for (uint64_t numChunks = 1u; numChunks <= uppBnd; ++numChunks) { + uint64_t depth = ceil_div(depthNumerator * numChunks, numCores); + uint64_t tokens = ceil_div(seqLength, numChunks); + uint64_t work = tokens + width; + double score = static_cast(depth) * static_cast(work); + if (score < bestScore) { + bestScore = score; + seqChunks = numChunks; } - int64_t idealBlockSize = CeilDiv(cuSeqlen, tokenCoreBudget); - if (idealBlockSize <= 0) { - idealBlockSize = 1; - } - td.tokenBlockSize = idealBlockSize; - td.tokenBlockCnt = CeilDiv(cuSeqlen, td.tokenBlockSize); } - td.hasExplicitTokenSeqRanges = 0; - td.explicitTokenSeqRangeCount = 0; -} - -int64_t ComputeWorkspaceSize(int32_t blockDim, int64_t batch, int64_t width, int64_t dim, bool hasInitialState) -{ - if (!hasInitialState) { - return 0; - } - constexpr int64_t kDtypeSize = 2; - constexpr int64_t kSyncBytesPerBlock = 32; - int64_t historyCount = (width - 1 > 0) ? width - 1 : 0; - return ASCENDC_RESERVED_WORKSPACE + static_cast(blockDim) * kSyncBytesPerBlock + - batch * historyCount * dim * kDtypeSize; + return {static_cast(channelsPerTile), static_cast(seqChunks)}; } } // namespace -HOST_API at::Tensor causal_conv1d_impl(const at::Tensor &x, const at::Tensor &weight, const at::Tensor &bias, - const at::Tensor &conv_states, const at::Tensor &query_start_loc, - const at::Tensor &cache_indices, const at::Tensor &has_initial_state, - const at::Tensor &num_accepted_tokens, int64_t activation_mode, - int64_t pad_slot_id, int64_t run_mode) +HOST_API at::Tensor causal_conv1d_impl(const at::Tensor &x, const at::Tensor &weight, const at::Tensor &conv_states, + const at::Tensor &query_start_loc, const at::Tensor &cache_indices, + const at::Tensor &has_initial_state, const at::Tensor &bias, + bool activation_mode, int64_t pad_slot_id) { - TORCH_CHECK(x.defined(), "x tensor must be defined"); - TORCH_CHECK(weight.defined(), "weight tensor must be defined"); - TORCH_CHECK(conv_states.defined(), "conv_states tensor must be defined"); - - TORCH_CHECK(x.dim() == 2 || x.dim() == 3, "x must be 2D or 3D tensor"); - TORCH_CHECK(weight.dim() == 2, "weight must be 2D tensor"); - + TORCH_CHECK(x.dim() == 2 || x.dim() == 3, "x must be 2D [cu_seqlen, dim] or 3D [batch, seq_len, dim]"); + TORCH_CHECK(weight.dim() == 2, "weight must be 2D [width, dim], got shape ", weight.sizes()); + const uint32_t width = static_cast(weight.size(0)); + TORCH_CHECK(isSupportedWidth(width), "Only filter widths 2..64 are supported, got ", weight.size(0)); + TORCH_CHECK(conv_states.dim() == 3, "conv_states must be 3D [num_cache_lines, state_len, dim]"); const at::ScalarType dtype = x.scalar_type(); - TORCH_CHECK(dtype == at::kBFloat16 || dtype == at::kHalf, "Only BF16 and FP16 are supported"); + TORCH_CHECK(dtype == at::kHalf || dtype == at::kBFloat16, "Only BF16 and FP16 are supported, got ", dtype); TORCH_CHECK(weight.scalar_type() == dtype, "weight dtype must match x dtype"); TORCH_CHECK(conv_states.scalar_type() == dtype, "conv_states dtype must match x dtype"); - - TORCH_CHECK(x.is_contiguous(), "x must be contiguous"); - TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous"); - TORCH_CHECK(conv_states.is_contiguous(), "conv_states must be contiguous"); - - int64_t dim = (x.dim() == 2) ? x.size(1) : x.size(2); - int64_t width = weight.size(0); - TORCH_CHECK(width >= MIN_WIDTH && width <= MAX_WIDTH, "Only support width in [2,4]"); - - int64_t inputMode = (x.dim() == 2) ? 0 : 1; - int64_t seqLen = (inputMode == 1) ? x.size(1) : 0; - int64_t batch = (inputMode == 1) ? x.size(0) : 0; - int64_t cuSeqlen = (inputMode == 0) ? x.size(0) : batch * seqLen; - - if (inputMode == 0) { - int64_t qslSize = query_start_loc.size(0); - TORCH_CHECK(qslSize >= 2, "query_start_loc must have at least 2 elements"); - batch = qslSize - 1; - } - - int64_t numCacheLines = conv_states.size(0); - int64_t stateLen = conv_states.size(1); - - int64_t activationInt = activation_mode; - - bool hasBias = bias.defined() && bias.numel() > 0; - bool hasCacheIndices = cache_indices.defined() && cache_indices.numel() > 0; - bool hasInitialState = has_initial_state.defined() && has_initial_state.numel() > 0; - bool hasNumAccept = num_accepted_tokens.defined() && num_accepted_tokens.numel() > 0; - bool isBf16 = (dtype == at::kBFloat16); - - at::Tensor y = at::empty_like(x); - - at::Tensor bias_tensor = hasBias ? bias : at::empty({0}, x.options()); - at::Tensor query_start_loc_tensor = (query_start_loc.defined() && query_start_loc.numel() > 0) - ? query_start_loc.to(at::kLong) - : at::empty({0}, x.options().dtype(at::kLong)); - at::Tensor cache_indices_tensor = - hasCacheIndices ? cache_indices.to(at::kLong) : at::empty({0}, x.options().dtype(at::kLong)); - at::Tensor has_initial_state_tensor = - hasInitialState ? has_initial_state.to(at::kLong) : at::empty({0}, x.options().dtype(at::kLong)); - at::Tensor num_accepted_tokens_tensor; - if (hasNumAccept) { - num_accepted_tokens_tensor = num_accepted_tokens.to(at::kInt); + TORCH_CHECK(query_start_loc.scalar_type() == at::kInt, "query_start_loc dtype must be int32"); + TORCH_CHECK(cache_indices.scalar_type() == at::kInt, "cache_indices dtype must be int32"); + TORCH_CHECK(has_initial_state.scalar_type() == at::kBool, "has_initial_state dtype must be bool"); + TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() && conv_states.is_contiguous(), + "inputs must be contiguous"); + + const bool has_bias = bias.numel() > 0; + uint32_t inputMode, batch, seqLen, dim; + if (x.dim() == 2) { + inputMode = 0; + dim = static_cast(x.size(1)); + seqLen = 0; + // Guard the size(0)-1 below: an empty/too-short qsl would underflow batch to ~4e9. + TORCH_CHECK(query_start_loc.dim() == 1 && query_start_loc.size(0) >= 2, + "query_start_loc must be 1D and have at least 2 elements"); + batch = static_cast(query_start_loc.size(0) - 1); } else { - num_accepted_tokens_tensor = at::empty({0}, x.options().dtype(at::kInt)); + inputMode = 1; + batch = static_cast(x.size(0)); + seqLen = static_cast(x.size(1)); + dim = static_cast(x.size(2)); } - - auto ascendc_platform = platform_ascendc::PlatformAscendCManager::GetInstance(); - int32_t maxAivCore = static_cast(ascendc_platform->GetCoreNumAiv()); - - CausalConv1dTilingData tilingData; - ComputeTilingData(dim, cuSeqlen, seqLen, batch, inputMode, width, stateLen, numCacheLines, activationInt, - pad_slot_id, hasBias, hasCacheIndices, hasInitialState, hasNumAccept, isBf16, maxAivCore, - run_mode, tilingData); - - int64_t totalBlocks = (run_mode == CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE) - ? (tilingData.batch * tilingData.baseDimCnt) - : (tilingData.tokenBlockCnt * tilingData.baseDimCnt); - int32_t blockDim = std::min(maxAivCore, static_cast(totalBlocks)); - if (blockDim <= 0) { - blockDim = 1; - } - - int32_t libApiWorkspaceSize = static_cast(ascendc_platform->GetLibApiWorkSpaceSize()); - - int64_t ws = ComputeWorkspaceSize(blockDim, batch, width, dim, hasInitialState); - int64_t totalWorkspace = std::max(static_cast(libApiWorkspaceSize), ws); - if (totalWorkspace <= 0) { - totalWorkspace = libApiWorkspaceSize; + TORCH_CHECK(batch > 0 && dim > 0, "bad batch/dim"); + // cache_indices[seq] and has_initial_state[seq] are read per sequence in both layouts. + TORCH_CHECK(cache_indices.dim() == 1 && cache_indices.size(0) >= static_cast(batch), + "cache_indices must be 1D and have size >= batch"); + TORCH_CHECK(has_initial_state.dim() == 1 && has_initial_state.size(0) >= static_cast(batch), + "has_initial_state must be 1D and have size >= batch"); + TORCH_CHECK(dim % 16 == 0, "dim must be multiple of 16 for fp16/bf16 alignment, but got ", dim); + TORCH_CHECK(weight.size(1) == static_cast(dim), "weight.shape[1] must equal dim"); + TORCH_CHECK(conv_states.size(2) == static_cast(dim), "conv_states.shape[2] must equal dim"); + if (has_bias) { // bias is read in the I/O dtype and cast to fp32 in the kernel + TORCH_CHECK(bias.dim() == 1 && bias.size(0) == static_cast(dim), "bias must be 1D [dim]"); + TORCH_CHECK(bias.scalar_type() == dtype, "bias dtype must match x dtype"); + TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous"); } + const uint32_t stateLen = static_cast(conv_states.size(1)); + TORCH_CHECK(stateLen >= width - 1, "state_len must be >= width-1"); + + auto plat = platform_ascendc::PlatformAscendCManager::GetInstance(); + TORCH_CHECK(plat != nullptr, "no AscendC platform"); + const uint32_t core_num = static_cast(plat->GetCoreNumAiv()); + TORCH_CHECK(core_num > 0, "bad core_num"); + + // ---- launch grid: (channel tile) x (sequence chunk) work units, sized to + // fill all AIV cores. Batch parallelism comes first; if batch alone can't fill + // the cores we split the channel axis into tiles and the L axis into chunks. ---- + const uint32_t avgSeqLen = + (inputMode == 1) ? seqLen : std::max(1u, static_cast(x.size(0)) / batch); + const uint32_t ringSize = roundUpToPow2(width); // compile-time ring variant to launch + + uint64_t ubBytes = 0; + plat->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubBytes); + constexpr uint64_t WIDE_UB_MIN_BYTES = 248u * 1024u; // UB the wide table is sized for + const bool wideUb = (ubBytes >= WIDE_UB_MIN_BYTES); + const uint32_t maxChannelsPerTile = maxTileWidthForRing(ringSize, wideUb); // UB-bound tile width + + const auto [channelsPerTile, seqChunks] = + tiling_causal_conv1d(core_num, batch, dim, avgSeqLen, width, maxChannelsPerTile); + const uint32_t channelTiles = ceil_div(dim, channelsPerTile); + + // weight/bias enter in the I/O dtype (fp16/bf16) and are cast to fp32 inside the + // kernel; pass a native empty placeholder when there is no bias. + const at::Tensor biasArg = has_bias ? bias : at::empty({0}, x.options()); + at::Tensor y = at::empty_like(x); - int32_t tilingSize = - (static_cast(sizeof(CausalConv1dTilingData)) + PADDING_BYTE - 1) / PADDING_BYTE * PADDING_BYTE; - - CausalConv1dTilingKey key{dim, - cuSeqlen, - seqLen, - batch, - inputMode, - width, - stateLen, - numCacheLines, - activationInt, - pad_slot_id, - run_mode, - hasBias ? 1 : 0, - hasCacheIndices ? 1 : 0, - hasInitialState ? 1 : 0, - hasNumAccept ? 1 : 0}; - uint64_t hashValue = CausalConv1dTilingKeyHash{}(key); - - static auto globalTilingBuffer = at::empty({tilingSize * static_cast(MAX_CAPTURE_NUM)}, - at::TensorOptions().dtype(at::kByte).device(x.options().device())); - - auto copyTilingToDevice = [&]() { - auto cpuTiling = at::empty({tilingSize}, at::kByte); - std::memcpy(cpuTiling.data_ptr(), &tilingData, sizeof(CausalConv1dTilingData)); - return TorchNpuHelper::CopyTensorHostToDevice(cpuTiling); - }; - - at::Tensor tilingTensor; - if (g_causalConv1dCaptureMap.find(hashValue) != g_causalConv1dCaptureMap.end()) { - tilingTensor = - at::from_blob(globalTilingBuffer.data_ptr() + (tilingSize * g_causalConv1dCaptureMap[hashValue]), - tilingSize, at::kByte); - } else if (g_causalConv1dCaptureNum >= MAX_CAPTURE_NUM) { - tilingTensor = copyTilingToDevice(); + // conv has one task per (batch, channel tile, seq chunk); the writeback only + // touches the sequence tail, so it drops the seq-chunk axis. Cap block dim at the cores. + const uint32_t blockDimConv = std::min(batch * channelTiles * seqChunks, core_num); + const uint32_t blockDimWb = std::min(batch * channelTiles, core_num); + const uint32_t actFlag = activation_mode ? 1u : 0u; + const uint32_t biasFlag = has_bias ? 1u : 0u; + const int32_t padSlot = static_cast(pad_slot_id); + + // Launch the ring = roundUpToPow2(width) variant (entry suffix rs) and pass + // the actual width as the runtime K. launch(suffix) fires the conv + writeback pair; + // the per-dtype switch over the ring-size list selects the entry. +#define launch(suffix) \ + do { \ + EXEC_KERNEL_CMD(causal_conv1d_##suffix, blockDimConv, x, weight, biasArg, conv_states, query_start_loc, \ + cache_indices, has_initial_state, y, dim, batch, inputMode, seqLen, stateLen, width, \ + channelsPerTile, channelTiles, seqChunks, actFlag, biasFlag, padSlot); \ + EXEC_KERNEL_CMD(causal_conv1d_wb_##suffix, blockDimWb, x, conv_states, query_start_loc, cache_indices, \ + has_initial_state, dim, batch, inputMode, seqLen, stateLen, width, channelsPerTile, \ + channelTiles, padSlot); \ + } while (0) +#define DISPATCH_HALF(ringSize, maxTileWidth) \ + case ringSize: \ + launch(rs##ringSize##_half); \ + break; +#define DISPATCH_BF16(ringSize, maxTileWidth) \ + case ringSize: \ + launch(rs##ringSize##_bf16); \ + break; + if (dtype == at::kHalf) { + switch (ringSize) { + FOR_EACH_RING_SIZE(DISPATCH_HALF) + default: + break; + } } else { - g_causalConv1dCaptureMap[hashValue] = g_causalConv1dCaptureNum; - auto deviceTiling = copyTilingToDevice(); - globalTilingBuffer - .slice(0, g_causalConv1dCaptureNum * tilingSize, g_causalConv1dCaptureNum * tilingSize + tilingSize) - .copy_(deviceTiling); - g_causalConv1dCaptureNum++; - tilingTensor = - at::from_blob(globalTilingBuffer.data_ptr() + (tilingSize * g_causalConv1dCaptureMap[hashValue]), - tilingSize, at::kByte); + switch (ringSize) { + FOR_EACH_RING_SIZE(DISPATCH_BF16) + default: + break; + } } - - auto workspaceTensor = - at::empty({totalWorkspace}, at::TensorOptions().dtype(at::kByte).device(x.options().device())); - - EXEC_KERNEL_CMD(causal_conv1d, blockDim, x, weight, conv_states, bias_tensor, query_start_loc_tensor, - cache_indices_tensor, has_initial_state_tensor, num_accepted_tokens_tensor, y, workspaceTensor, - tilingTensor); - +#undef DISPATCH_BF16 +#undef DISPATCH_HALF +#undef launch return y; } diff --git a/csrc/causal_conv1d/op_host/causal_conv1d.h b/csrc/causal_conv1d/op_host/causal_conv1d.h index a39a055b4..2ed6f7d76 100644 --- a/csrc/causal_conv1d/op_host/causal_conv1d.h +++ b/csrc/causal_conv1d/op_host/causal_conv1d.h @@ -1,35 +1,23 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR - * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See - * LICENSE in the root of the software repository for the full text of the License. - */ - /*! * \file causal_conv1d.h - * \brief causal_conv1d host-side function declaration + * \brief host-side declaration for the PTO-ISA causal_conv1d (drop-in alternative). */ - -#ifndef CUSTOM_CAUSAL_CONV1D_HOST_H_ -#define CUSTOM_CAUSAL_CONV1D_HOST_H_ +#ifndef CAUSAL_CONV1D_HOST_H_ +#define CAUSAL_CONV1D_HOST_H_ #include + #include "defines.h" namespace sglang { namespace npu_kernel { -HOST_API at::Tensor causal_conv1d_impl(const at::Tensor &x, const at::Tensor &weight, const at::Tensor &bias, - const at::Tensor &conv_states, const at::Tensor &query_start_loc, - const at::Tensor &cache_indices, const at::Tensor &has_initial_state, - const at::Tensor &num_accepted_tokens, int64_t activation_mode, - int64_t pad_slot_id, int64_t run_mode); +HOST_API at::Tensor causal_conv1d_impl(const at::Tensor &x, const at::Tensor &weight, const at::Tensor &conv_states, + const at::Tensor &query_start_loc, const at::Tensor &cache_indices, + const at::Tensor &has_initial_state, const at::Tensor &bias, + bool activation_mode, int64_t pad_slot_id); } // namespace npu_kernel } // namespace sglang -#endif // CUSTOM_CAUSAL_CONV1D_HOST_H_ +#endif // CAUSAL_CONV1D_HOST_H_ diff --git a/csrc/causal_conv1d/op_host/stub/aclrtlaunch_causal_conv1d.h b/csrc/causal_conv1d/op_host/stub/aclrtlaunch_causal_conv1d.h deleted file mode 100644 index b75f4e7f7..000000000 --- a/csrc/causal_conv1d/op_host/stub/aclrtlaunch_causal_conv1d.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef HEADER_ACLRTLAUNCH_CUSTOM_CAUSAL_CONV1D_H -#define HEADER_ACLRTLAUNCH_CUSTOM_CAUSAL_CONV1D_H -#include "acl/acl_base.h" - -#ifndef ACLRT_LAUNCH_KERNEL -#define ACLRT_LAUNCH_KERNEL(kernel_func) aclrtlaunch_##kernel_func -#endif - -extern "C" uint32_t aclrtlaunch_causal_conv1d(uint32_t numBlocks, aclrtStream stream, void *x, void *weight, - void *convStates, void *bias, void *queryStartLoc, void *cacheIndices, - void *initialStateMode, void *numAcceptedTokens, void *y, void *workspace, - void *tiling); -#endif diff --git a/csrc/causal_conv1d/op_kernel/arch35/causal_conv1d_regbase.h b/csrc/causal_conv1d/op_kernel/arch35/causal_conv1d_regbase.h deleted file mode 100644 index fb6aa0886..000000000 --- a/csrc/causal_conv1d/op_kernel/arch35/causal_conv1d_regbase.h +++ /dev/null @@ -1,183 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -/*! - * \file causal_conv1d_regbase.h - * \brief - */ -#ifndef CUSTOM_CAUSAL_CONV1D_REGBASE_H -#define CUSTOM_CAUSAL_CONV1D_REGBASE_H - -namespace NsCausalConv1d { -using namespace AscendC; -using namespace AscendC::MicroAPI; - -constexpr uint16_t V_LENGTH = VECTOR_REG_WIDTH / sizeof(float); - -constexpr CastTrait castTraitB16ToB32 = {RegLayout::ZERO, SatMode::UNKNOWN, MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; - -template -__aicore__ inline void ComputeFnRollingOutputRegbase(LocalTensor ring, LocalTensor currF, - LocalTensor state0F, LocalTensor weightF, - uint32_t dataCount) -{ - __ubuf__ T *ringAddr = (__ubuf__ T *)ring.GetPhyAddr(); - __ubuf__ float *currFAddr = (__ubuf__ float *)currF.GetPhyAddr(); - __ubuf__ float *state0FAddr = (__ubuf__ float *)state0F.GetPhyAddr(); - __ubuf__ float *weightFAddr = (__ubuf__ float *)weightF.GetPhyAddr(); - - uint16_t colLoopTimes = static_cast(Ceil(dataCount, V_LENGTH)); - __VEC_SCOPE__ - { - RegTensor ring; - RegTensor currF; - RegTensor state0F; - RegTensor weightF; - RegTensor tmp; - MaskReg pregLoop; - for (uint16_t j = 0; j < colLoopTimes; j++) { - pregLoop = UpdateMask(dataCount); - DataCopy(ring, ringAddr + j * V_LENGTH); - DataCopy(state0F, state0FAddr + j * V_LENGTH); - DataCopy(weightF, weightFAddr + j * V_LENGTH); - Cast(currF, ring, pregLoop); - Mul(currF, currF, weightF, pregLoop); - Add(state0F, state0F, currF, pregLoop); - if constexpr (hasActivation) { - Muls(tmp, state0F, -1.0f, pregLoop); - Exp(tmp, tmp, pregLoop); - Adds(tmp, tmp, 1.0f, pregLoop); - Div(currF, state0F, tmp, pregLoop); - DataCopy(currFAddr + j * V_LENGTH, currF, pregLoop); - } else { - DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); - } - } - } -} - -template -static __simd_vf__ inline void AdvanceFnLocalPartialsWidthTwo(__ubuf__ T *ringAddr, __ubuf__ float *weight0FAddr, - __ubuf__ float *state0FAddr, uint32_t dataCount, - uint16_t colLoopTimes) -{ - RegTensor ring; - RegTensor currF; - RegTensor weight0F; - RegTensor state0F; - MaskReg pregLoop; - for (uint16_t j = 0; j < colLoopTimes; j++) { - pregLoop = UpdateMask(dataCount); - DataCopy(ring, ringAddr + j * V_LENGTH); - DataCopy(weight0F, weight0FAddr + j * V_LENGTH); - Cast(currF, ring, pregLoop); - Mul(state0F, currF, weight0F, pregLoop); - DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); - } -} - -template -static __simd_vf__ inline void AdvanceFnLocalPartialsWidthThree(__ubuf__ T *ringAddr, __ubuf__ float *weight0FAddr, - __ubuf__ float *weight1FAddr, - __ubuf__ float *state0FAddr, - __ubuf__ float *state1FAddr, uint32_t dataCount, - uint16_t colLoopTimes) -{ - RegTensor ring; - RegTensor currF; - RegTensor weight0F; - RegTensor weight1F; - RegTensor state0F; - RegTensor state1F; - MaskReg pregLoop; - for (uint16_t j = 0; j < colLoopTimes; j++) { - pregLoop = UpdateMask(dataCount); - DataCopy(ring, ringAddr + j * V_LENGTH); - DataCopy(state1F, state1FAddr + j * V_LENGTH); - Cast(currF, ring, pregLoop); - DataCopy(weight1F, weight1FAddr + j * V_LENGTH); - Mul(state0F, currF, weight1F, pregLoop); - DataCopy(weight0F, weight0FAddr + j * V_LENGTH); - Add(state0F, state0F, state1F, pregLoop); - Mul(state1F, currF, weight0F, pregLoop); - DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); - DataCopy(state1FAddr + j * V_LENGTH, state1F, pregLoop); - } -} - -template -static __simd_vf__ inline void -AdvanceFnLocalPartialsWidthFour(__ubuf__ T *ringAddr, __ubuf__ float *weight0FAddr, __ubuf__ float *weight1FAddr, - __ubuf__ float *weight2FAddr, __ubuf__ float *state0FAddr, __ubuf__ float *state1FAddr, - __ubuf__ float *state2FAddr, uint32_t dataCount, uint16_t colLoopTimes) -{ - RegTensor ring; - RegTensor currF; - RegTensor weight0F; - RegTensor weight1F; - RegTensor weight2F; - RegTensor state0F; - RegTensor state1F; - RegTensor state2F; - MaskReg pregLoop; - for (uint16_t j = 0; j < colLoopTimes; j++) { - pregLoop = UpdateMask(dataCount); - DataCopy(ring, ringAddr + j * V_LENGTH); - DataCopy(state1F, state1FAddr + j * V_LENGTH); - DataCopy(state2F, state2FAddr + j * V_LENGTH); - Cast(currF, ring, pregLoop); - DataCopy(weight2F, weight2FAddr + j * V_LENGTH); - Mul(state0F, currF, weight2F, pregLoop); - DataCopy(weight1F, weight1FAddr + j * V_LENGTH); - Add(state0F, state0F, state1F, pregLoop); - Mul(state1F, currF, weight1F, pregLoop); - DataCopy(weight0F, weight0FAddr + j * V_LENGTH); - Add(state1F, state1F, state2F, pregLoop); - Mul(state2F, currF, weight0F, pregLoop); - DataCopy(state0FAddr + j * V_LENGTH, state0F, pregLoop); - DataCopy(state1FAddr + j * V_LENGTH, state1F, pregLoop); - DataCopy(state2FAddr + j * V_LENGTH, state2F, pregLoop); - } -} - -template -__aicore__ inline void AdvanceFnLocalPartialsRegbase(LocalTensor ring, LocalTensor weightF, - LocalTensor state0F, LocalTensor state1F, - LocalTensor state2F, uint32_t dataCount, - uint32_t weightStep) -{ - uint16_t colLoopTimes = static_cast(Ceil(dataCount, V_LENGTH)); - - __ubuf__ T *ringAddr = (__ubuf__ T *)ring.GetPhyAddr(); - __ubuf__ float *weight0FAddr = (__ubuf__ float *)weightF.GetPhyAddr(); - __ubuf__ float *state0FAddr = (__ubuf__ float *)state0F.GetPhyAddr(); - if constexpr (kTemplateWidth == 2) { - AscendC::VF_CALL>(ringAddr, weight0FAddr, state0FAddr, dataCount, - colLoopTimes); - } else if constexpr (kTemplateWidth == 3) { - __ubuf__ float *weight1FAddr = weight0FAddr + weightStep; - __ubuf__ float *state1FAddr = (__ubuf__ float *)state1F.GetPhyAddr(); - AscendC::VF_CALL>(ringAddr, weight0FAddr, weight1FAddr, state0FAddr, - state1FAddr, dataCount, colLoopTimes); - } else if constexpr (kTemplateWidth == 4) { - __ubuf__ float *weight1FAddr = weight0FAddr + weightStep; - __ubuf__ float *weight2FAddr = weight1FAddr + weightStep; - __ubuf__ float *state1FAddr = (__ubuf__ float *)state1F.GetPhyAddr(); - __ubuf__ float *state2FAddr = (__ubuf__ float *)state2F.GetPhyAddr(); - AscendC::VF_CALL>(ringAddr, weight0FAddr, weight1FAddr, weight2FAddr, - state0FAddr, state1FAddr, state2FAddr, dataCount, - colLoopTimes); - } -} - -} // namespace NsCausalConv1d - -#endif // CUSTOM_CAUSAL_CONV1D_REGBASE_H diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d.cpp b/csrc/causal_conv1d/op_kernel/causal_conv1d.cpp index 26cffcd60..7b1a2e26c 100644 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d.cpp +++ b/csrc/causal_conv1d/op_kernel/causal_conv1d.cpp @@ -1,104 +1,516 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -/*! - * \file causal_conv1d.cpp - * \brief causal_conv1d kernel entry with runtime dispatch - */ - -#include "causal_conv1d_fn.h" -#include "causal_conv1d_update.h" - -using namespace AscendC; -using namespace NsCausalConv1d; - -namespace { - -template -__aicore__ inline void DispatchFn(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, GM_ADDR queryStartLoc, - GM_ADDR cacheIndices, GM_ADDR initialStateMode, GM_ADDR numAcceptedTokens, GM_ADDR y, - GM_ADDR workspace, const __gm__ CausalConv1dTilingData *tilingData, uint32_t widthKey, - uint32_t fnPlanKey) +// causal_conv1d.cpp — PTO-ISA depthwise causal conv1d + bias + (opt) SiLU. +// +// Drop-in for sgl's AscendC causal_conv1d, written on the PTO tile ISA (like +// csrc/mega_chunk_gdn). Same op: +// y[t,c] = act( bias[c] + sum_{k=0..K-1} W[k,c]*xext[t+k,c] ), +// xext = [history(K-1 rows, from conv_states if has_initial_state else 0), x] +// Per channel a K-tap depthwise filter. The filter width K is a RUNTIME argument; +// the compile-time template parameter is the accumulator-ring size RS (a power of +// two, K <= RS). The host routes a request to the RS = roundUpToPow2(K) variant, so +// any width in [2, 64] is served by the six compiled RS variants {2,4,8,16,32,64}. +// fp16/bf16 I/O, fp32 accumulate. Weights/bias enter native and are cast to fp32 on device. +// +// Work grid (in-kernel, uniform task striding): batch x blocksPerSeq x lchunks. +// blocksPerSeq = ceil(dim/col_w) (channel tiles), lchunks splits the L axis so +// all cores stay busy even at small batch. Causal halo = K-1 replayed rows per +// chunk. State writeback is a SEPARATE launch (causal_conv1d_wb_*) to avoid a race +// between chunk-0's history read and the tail chunk's state write. +// +// Scalars (qsl/cidx/hinit) are read by direct __gm__ pointer indexing (same as +// mega_kernel's cu_seqlens). NOTE codegen needs "type* name", not "type *name". + +#include +#if defined(__NPU_ARCH__) // device compile pass only; PTO arch-aware buffer sizes +#include +#include "kernel_operator.h" // AscendC::PipeBarrier -- arch-portable vector barrier +#endif + +// In-core vector barrier. A5 (dav-c310) has no PIPE_V barrier in the VF/RegBase model, so +// the legacy pipe_barrier(PIPE_V) intrinsic cannot be used unconditionally. AscendC's +// templated barrier resolves per architecture, which avoids hardcoding an arch check that +// would silently become a no-op on a future arch that does need it. +#define PIPE_BARRIER_VEC() AscendC::PipeBarrier() + +// clang-format off +// The AscendC launch codegen parses the expanded __global__ signature and needs +// "type* name" (pointer glued to the type); PointerAlignment: Right would turn +// it into "type *name" and the generated launch stub loses the parameter names. +#ifndef GM_ADDR +#define GM_ADDR __gm__ uint8_t* +#endif +// clang-format on + +using namespace pto; + +namespace cc1d { + +// RS (compile-time, power of two) sizes the accumulator ring and the entire UB +// layout; K (runtime, <= RS) only drives loop bounds, so one RS variant serves +// every width with roundUpToPow2(width) == RS. MAX_W is the compile-time per-RS +// channel-tile capacity. A5 (dav-c310) has a larger UB than A2/A3; the static_assert +// in convChunk checks the chosen (RS, MAX_W) fits. +#ifdef __DAV_C310__ +constexpr uint32_t UB_BYTES_PER_CORE = 248u * 1024u; +#else +constexpr uint32_t UB_BYTES_PER_CORE = 192u * 1024u; +#endif + +template +AICORE inline void applySiluToTile(TileT &dst, TileT &src, TileT &tmp) { - if (fnPlanKey == CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS) { - if (widthKey == CAUSAL_CONV1D_TPL_WIDTH_2) { - RunCausalConv1dFn( - x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, - workspace, tilingData); - } else if (widthKey == CAUSAL_CONV1D_TPL_WIDTH_3) { - RunCausalConv1dFn( - x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, - workspace, tilingData); + using T = typename TileT::DType; + TMULS(tmp, src, (T)-1); + PIPE_BARRIER_VEC(); + TEXP(tmp, tmp); + PIPE_BARRIER_VEC(); + TADDS(tmp, tmp, (T)1); + PIPE_BARRIER_VEC(); + TDIV(dst, src, tmp); +} + +// One conv chunk: outputs [l0,l1) for channels [c0,c0+lanes) of a sequence whose +// tokens start at element row `start` (token index). history from convStates. +template +AICORE inline void convChunk(__gm__ IoElemType *x, __gm__ IoElemType *y, __gm__ IoElemType *wgt, __gm__ IoElemType *bia, + __gm__ IoElemType *convStates, uint32_t dim, uint32_t stateLen, uint32_t K, int32_t start, + int32_t cacheIdx, bool hasInit, uint32_t c0, int32_t lanes, int32_t l0, int32_t l1, + uint32_t activation, uint32_t hasBias) +{ + using GlobalShape = pto::Shape<1, 1, 1, 1, DYNAMIC>; + using GlobalStride = pto::Stride<1, 1, 1, 1, 1>; + using GlobalIoTensor = pto::GlobalTensor; + using IoTile = Tile; + using AccumTile = Tile; + + constexpr uint32_t accumTileBytes = MAX_W * sizeof(float); + constexpr uint32_t ioTileBytes = MAX_W * sizeof(IoElemType); + + // UB byte offsets, all compile-time on RS: the layout is sized for the worst + // case K == RS so no offset depends on the runtime K. fp32 region: RS weights + // (weight k at k*accumTileBytes) | bias | RS accumulators | RS-1 temps | xin_f. + // Then the I/O region: 4 ioTileBytes-sized tiles (input load double-buffered: + // xin_h[0] | out0 | out1 | xin_h[1]). + constexpr uint32_t ubBiasOffset = RS * accumTileBytes; + constexpr uint32_t ubAccumRingBase = (RS + 1u) * accumTileBytes; + constexpr uint32_t ubProductBase = (2u * RS + 1u) * accumTileBytes; // temp k at ubProductBase+(k-1)*accumTileBytes + constexpr uint32_t ubInputFp32 = (3u * RS) * accumTileBytes; + constexpr uint32_t ubIoBase = (3u * RS + 1u) * accumTileBytes; // I/O region base + static_assert(ubIoBase + 4u * ioTileBytes <= UB_BYTES_PER_CORE, + "conv1d UB exceeds UB_BYTES_PER_CORE: lower RS/MAX_W or raise it"); + // NOTE: keep these `const`, NOT `constexpr`. These arrays are indexed by a + // runtime value in the hot loop; making them constexpr makes the ascendc/cce + // compiler emit ~5x slower device code (measured 951us vs 186us at B8/L512/ + // d6144). (Harmless under the bisheng JIT path, but not here.) + const uint32_t ubOutputOffset[2] = {ubIoBase + ioTileBytes, ubIoBase + 2u * ioTileBytes}; + const uint32_t ubInputOffset[2] = {ubIoBase, ubIoBase + 3u * ioTileBytes}; + + // weights/bias arrive native (fp16/bf16); cast them to the resident fp32 tiles on + // device. The accumulator/temp/xin_f region is idle until the input loop, so use it + // as scratch: stage all K(+bias) native tiles there (the loads pipeline on MTE2 just + // like a plain load), one MTE2->V barrier, then cast each (the TCVTs pipeline on V). + // EVENT_ID3 here is the same load barrier the original used and is reused below. + constexpr uint32_t ubStageBase = ubAccumRingBase; // accumulators + temps + xin_f are free here + static_assert((RS + 1u) * ioTileBytes <= 2u * RS * accumTileBytes, + "conv1d: native weight/bias staging does not fit the scratch region"); + // The PREVIOUS task's output phase reads this same region with V (TCVT(outT, acc)); + // drain V before our staging MTE2 overwrites it -- otherwise a cross-task WAR + // corrupts that task's output. Self-contained on ID0 (clean here; reused by the loop). + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + for (uint32_t k = 0; k < K; ++k) { + GlobalIoTensor wG(wgt + (uint64_t)k * dim + c0, {lanes}); + IoTile wStage(lanes); + TASSIGN(wStage, ubStageBase + k * ioTileBytes); + TLOAD(wStage, wG); + } + if (hasBias) { + GlobalIoTensor bG(bia + c0, {lanes}); + IoTile bStage(lanes); + TASSIGN(bStage, ubStageBase + K * ioTileBytes); + TLOAD(bStage, bG); + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); // all native tiles staged before any cast + for (uint32_t k = 0; k < K; ++k) { + IoTile wStage(lanes); + AccumTile wT(lanes); + TASSIGN(wStage, ubStageBase + k * ioTileBytes); + TASSIGN(wT, k * accumTileBytes); + TCVT(wT, wStage, pto::RoundMode::CAST_NONE); + } + if (hasBias) { + IoTile bStage(lanes); + AccumTile bT(lanes); + TASSIGN(bStage, ubStageBase + K * ioTileBytes); + TASSIGN(bT, ubBiasOffset); + TCVT(bT, bStage, pto::RoundMode::CAST_NONE); + } + // The cast TCVTs (V) finish before the input loop's first TMUL/TCVT (also V, in + // program order) reuses this scratch region -- no extra sync needed. + + // double-buffered input: two load slots with independent handshakes. + // EVENT_ID3 is reused here (the weight/bias load above already consumed it). + const event_t IEV[2] = {EVENT_ID0, EVENT_ID3}; + set_flag(PIPE_V, PIPE_MTE2, IEV[0]); // xin_h[0] initially free + set_flag(PIPE_V, PIPE_MTE2, IEV[1]); // xin_h[1] initially free + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); + + // first input row to process (signed): l0==0 with history -> replay K-1 rows. + const int32_t halo = (int32_t)K - 1; // K-1 as signed (history rows go negative) + const bool zeroPad = (l0 == 0) && !hasInit; + int32_t jstart; + if (l0 == 0) + jstart = hasInit ? -halo : 0; + else + jstart = l0 - halo; + + // PROLOGUE: load the first input row (jstart) so iter 0 can prefetch the next. + // Input row index e: e>=0 -> x[start+e]; e<0 -> conv_states history row (K-1)+e. + if (jstart < l1) { + IoTile xin_h0(lanes); + TASSIGN(xin_h0, ubInputOffset[0]); + wait_flag(PIPE_V, PIPE_MTE2, IEV[0]); + if (jstart >= 0) { + GlobalIoTensor xG(x + (uint64_t)(start + jstart) * dim + c0, {lanes}); + TLOAD(xin_h0, xG); } else { - RunCausalConv1dFn( - x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, - workspace, tilingData); + const int32_t hi = halo + jstart; + GlobalIoTensor hG(convStates + ((uint64_t)cacheIdx * stateLen + hi) * dim + c0, {lanes}); + TLOAD(xin_h0, hG); } - } else { - if (widthKey == CAUSAL_CONV1D_TPL_WIDTH_2) { - RunCausalConv1dFn( - x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, - workspace, tilingData); - } else if (widthKey == CAUSAL_CONV1D_TPL_WIDTH_3) { - RunCausalConv1dFn( - x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, - workspace, tilingData); - } else { - RunCausalConv1dFn( - x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, - workspace, tilingData); + set_flag(PIPE_MTE2, PIPE_V, IEV[0]); + } + + for (int32_t j = jstart; j < l1; ++j) { + const uint32_t par = (j - jstart) & 1u; + IoTile xin_h(lanes); + AccumTile xin_f(lanes); + TASSIGN(xin_h, ubInputOffset[par]); + TASSIGN(xin_f, ubInputFp32); + + // (1) consume current row (loaded by prologue / previous prefetch) in buffer par + wait_flag(PIPE_MTE2, PIPE_V, IEV[par]); + TCVT(xin_f, xin_h, pto::RoundMode::CAST_NONE); + set_flag(PIPE_V, PIPE_MTE2, IEV[par]); + + // (2) prefetch next row (x or conv_states history) into the OTHER buffer + if (j + 1 < l1) { + const int32_t e = j + 1; + const uint32_t p1 = par ^ 1u; + IoTile xin_hn(lanes); + TASSIGN(xin_hn, ubInputOffset[p1]); + wait_flag(PIPE_V, PIPE_MTE2, IEV[p1]); + if (e >= 0) { + GlobalIoTensor xG(x + (uint64_t)(start + e) * dim + c0, {lanes}); + TLOAD(xin_hn, xG); + } else { + const int32_t hi = halo + e; + GlobalIoTensor hG(convStates + ((uint64_t)cacheIdx * stateLen + hi) * dim + c0, {lanes}); + TLOAD(xin_hn, hG); + } + set_flag(PIPE_MTE2, PIPE_V, IEV[p1]); + } + + PIPE_BARRIER_VEC(); + + const bool startAll = zeroPad && (j == 0); +#ifdef __DAV_C310__ + for (uint32_t k = 0; k < K; ++k) { + const int32_t out = j + halo - (int32_t)k; + if (out < l0 || out >= l1) continue; + AccumTile wT(lanes); + AccumTile acc(lanes); + TASSIGN(wT, k * accumTileBytes); + TASSIGN(acc, ubAccumRingBase + (out & (RS - 1u)) * accumTileBytes); + if (startAll || k == 0) { + TMUL(acc, xin_f, wT); + } else { + TMULADDDST(acc, xin_f, wT); + } + } + PIPE_BARRIER_VEC(); +#else + for (uint32_t k = 0; k < K; ++k) { + const int32_t out = j + halo - (int32_t)k; + if (out < l0 || out >= l1) continue; + AccumTile wT(lanes); + TASSIGN(wT, k * accumTileBytes); + if (startAll || k == 0) { + AccumTile acc(lanes); + TASSIGN(acc, ubAccumRingBase + (out & (RS - 1u)) * accumTileBytes); + TMUL(acc, xin_f, wT); + } else { + AccumTile t(lanes); + TASSIGN(t, ubProductBase + (k - 1u) * accumTileBytes); + TMUL(t, xin_f, wT); + } } + PIPE_BARRIER_VEC(); + if (!startAll) { + for (uint32_t k = 1; k < K; ++k) { + const int32_t out = j + halo - (int32_t)k; + if (out < l0 || out >= l1) continue; + AccumTile acc(lanes); + AccumTile t(lanes); + TASSIGN(acc, ubAccumRingBase + (out & (RS - 1u)) * accumTileBytes); + TASSIGN(t, ubProductBase + (k - 1u) * accumTileBytes); + TADD(acc, acc, t); + } + } + PIPE_BARRIER_VEC(); +#endif + + if (j < l0) continue; // halo row + + const uint32_t slot = j & (RS - 1u); + const uint32_t ob = j & 1u; + const event_t oev = (event_t)(1u + ob); + AccumTile acc(lanes); + AccumTile tmp(lanes); + IoTile outT(lanes); + TASSIGN(acc, ubAccumRingBase + slot * accumTileBytes); + TASSIGN(tmp, ubProductBase); + TASSIGN(outT, ubOutputOffset[ob]); + + if (hasBias) { + AccumTile bT(lanes); + TASSIGN(bT, ubBiasOffset); + TADD(acc, acc, bT); + PIPE_BARRIER_VEC(); + } + if (activation) { + applySiluToTile(acc, acc, tmp); + PIPE_BARRIER_VEC(); + } + wait_flag(PIPE_MTE3, PIPE_V, oev); + TCVT(outT, acc, pto::RoundMode::CAST_NONE); + GlobalIoTensor yG(y + (uint64_t)(start + j) * dim + c0, {lanes}); + set_flag(PIPE_V, PIPE_MTE3, oev); + wait_flag(PIPE_V, PIPE_MTE3, oev); + TSTORE(yG, outT); + set_flag(PIPE_MTE3, PIPE_V, oev); } -} -} // namespace + wait_flag(PIPE_V, PIPE_MTE2, IEV[0]); + wait_flag(PIPE_V, PIPE_MTE2, IEV[1]); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); +} -extern "C" __global__ __aicore__ void causal_conv1d(GM_ADDR x, GM_ADDR weight, GM_ADDR convStates, GM_ADDR bias, - GM_ADDR queryStartLoc, GM_ADDR cacheIndices, - GM_ADDR initialStateMode, GM_ADDR numAcceptedTokens, GM_ADDR y, - GM_ADDR workspace, GM_ADDR tiling) +template +AICORE void runConv(__gm__ IoElemType *x, __gm__ IoElemType *wgt, __gm__ IoElemType *bia, __gm__ IoElemType *convStates, + __gm__ int32_t *qsl, __gm__ int32_t *cidx, __gm__ uint8_t *hinit, __gm__ IoElemType *y, + uint32_t dim, uint32_t batch, uint32_t inputMode, uint32_t seqLen, uint32_t stateLen, uint32_t K, + uint32_t col_w, uint32_t blocksPerSeq, uint32_t lchunks, uint32_t activation, uint32_t hasBias, + int32_t padSlot) { - REGISTER_TILING_DEFAULT(CausalConv1dTilingData); - KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIV_1_0); - GM_ADDR userWorkspace = workspace; - if (workspace != nullptr) { - userWorkspace = AscendC::GetUserWorkspace(workspace); - } + set_mask_norm(); + set_vector_mask(-1, -1); + const uint32_t num_cores = get_block_num(); + const uint32_t core_id = get_block_idx(); + const uint32_t gridSize = batch * blocksPerSeq * lchunks; - auto tilingData = reinterpret_cast<__gm__ CausalConv1dTilingData *>(tiling); - auto runModeKey = static_cast(tilingData->runModeKey); - auto widthKey = static_cast(tilingData->widthKey); - auto fnPlanKey = static_cast(tilingData->fnPlanKey); - auto dtypeKey = static_cast(tilingData->dtypeKey); + for (uint32_t task = core_id; task < gridSize; task += num_cores) { + const uint32_t lc = task % lchunks; + const uint32_t t2 = task / lchunks; + const uint32_t db = t2 % blocksPerSeq; + const uint32_t seq = t2 / blocksPerSeq; - if (runModeKey == CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE) { - if (dtypeKey == 0) { - RunCausalConv1dUpdate(x, weight, bias, convStates, queryStartLoc, cacheIndices, - initialStateMode, numAcceptedTokens, y, userWorkspace, tilingData); + int32_t start, len; + if (inputMode == 0u) { + start = qsl[seq]; + len = qsl[seq + 1] - start; } else { - RunCausalConv1dUpdate(x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, - numAcceptedTokens, y, userWorkspace, tilingData); + start = seq * seqLen; + len = seqLen; } - return; + if (len == 0) continue; + const int32_t ci = cidx[seq]; + if (ci == padSlot) continue; + const bool hasInit = hinit[seq] != 0; + + // Clamp so no non-first chunk starts inside the causal halo (l0 = lc*lc_len >= K-1, + // so its history is in x, not conv_states) -- guards short/varlen sub-halo chunks. + uint32_t lc_len = (len + lchunks - 1) / lchunks; + if (lc_len < K - 1u) lc_len = K - 1u; + const int32_t l0 = lc * lc_len; + if (l0 >= len) continue; + int32_t l1 = l0 + lc_len; + if (l1 > len) l1 = len; + + const uint32_t c0 = db * col_w; + const uint32_t rem = dim - c0; + const int32_t lanes = (int32_t)(rem > col_w ? col_w : rem); + + convChunk(x, y, wgt, bia, convStates, dim, stateLen, K, start, ci, hasInit, c0, lanes, + l0, l1, activation, hasBias); } +} + +// Writeback: convStates[ci, 0:K-1, c0:] = last K-1 rows of xext (x tail / old hist). +template +AICORE void runWriteback(__gm__ IoElemType *x, __gm__ IoElemType *convStates, __gm__ int32_t *qsl, __gm__ int32_t *cidx, + __gm__ uint8_t *hinit, uint32_t dim, uint32_t batch, uint32_t inputMode, uint32_t seqLen, + uint32_t stateLen, uint32_t K, uint32_t col_w, uint32_t blocksPerSeq, int32_t padSlot) +{ + using GlobalShape = pto::Shape<1, 1, 1, 1, DYNAMIC>; + using GlobalStride = pto::Stride<1, 1, 1, 1, 1>; + using GlobalIoTensor = pto::GlobalTensor; + using IoTile = Tile; + using AccumTile = Tile; + constexpr uint32_t ioTileBytes = MAX_W * sizeof(IoElemType); + // compile-time scratch offset past the RS-1 reserved row slots (sized for the + // worst case K==RS so it never depends on the runtime K). + constexpr uint32_t SCRATCH_F32 = (RS - 1u) * ioTileBytes; // fp32 scratch for zeroing - if (dtypeKey == 0) { - uint32_t effectiveFnPlan = - (fnPlanKey == CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD) ? fnPlanKey : CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS; - DispatchFn(x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, - numAcceptedTokens, y, userWorkspace, tilingData, widthKey, effectiveFnPlan); - } else { - uint32_t effectiveFnPlan = - (fnPlanKey == CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD) ? fnPlanKey : CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS; - DispatchFn(x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, - y, userWorkspace, tilingData, widthKey, effectiveFnPlan); + set_mask_norm(); + set_vector_mask(-1, -1); + const uint32_t num_cores = get_block_num(); + const uint32_t core_id = get_block_idx(); + const uint32_t gridSize = batch * blocksPerSeq; + + // This core runs several tasks (strided) reusing the same UB tiles, so each + // task's reload (MTE2) must wait for the previous task's store (MTE3) to finish + // reading them -- otherwise the store races the reload. Start with UB "free". + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + for (uint32_t task = core_id; task < gridSize; task += num_cores) { + const uint32_t db = task % blocksPerSeq; + const uint32_t seq = task / blocksPerSeq; + int32_t start, len; + if (inputMode == 0u) { + start = qsl[seq]; + len = qsl[seq + 1] - start; + } else { + start = seq * seqLen; + len = seqLen; + } + if (len == 0) continue; + const int32_t ci = cidx[seq]; + if (ci == padSlot) continue; + const bool hasInit = hinit[seq] != 0; + const uint32_t c0 = db * col_w; + const uint32_t rem = dim - c0; + const int32_t lanes = (int32_t)(rem > col_w ? col_w : rem); + const int32_t halo = (int32_t)K - 1; + + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); // previous task's store done -> UB free to reload + // Phase A (MTE2): load K-1 source rows = xext[len .. len+K-2]. + // xext index e: e>=K-1 -> x[e-(K-1)]; e history (convStates if hasInit + // else zero). For the zero case load x[start] (finite) as a placeholder and + // zero it in phase B (avoids multiplying uninitialised UB). + for (int32_t i = 0; i < halo; ++i) { + const int32_t e = len + i; + const int32_t xrow = e - halo; + IoTile row(lanes); + TASSIGN(row, i * ioTileBytes); + if (xrow >= 0) { + GlobalIoTensor sG(x + (uint64_t)(start + xrow) * dim + c0, {lanes}); + TLOAD(row, sG); + } else if (hasInit) { + GlobalIoTensor hG(convStates + ((uint64_t)ci * stateLen + e) * dim + c0, {lanes}); + TLOAD(row, hG); + } else { + GlobalIoTensor sG(x + (uint64_t)start * dim + c0, {lanes}); // placeholder (len>=1) + TLOAD(row, sG); + } + } + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + // Phase B (V): zero the pure-history rows when there is no initial state. + // (TMULS has no bf16 overload, so zero via an fp32 round-trip of the finite + // placeholder: IoElemType -> fp32 -> *0 -> IoElemType. Works for fp16 and bf16.) + for (int32_t i = 0; i < halo; ++i) { + const int32_t e = len + i; + const int32_t xrow = e - halo; + if (xrow < 0 && !hasInit) { + IoTile row(lanes); + AccumTile f32(lanes); + TASSIGN(row, i * ioTileBytes); + TASSIGN(f32, SCRATCH_F32); + TCVT(f32, row, pto::RoundMode::CAST_NONE); + PIPE_BARRIER_VEC(); + TMULS(f32, f32, 0.0f); + PIPE_BARRIER_VEC(); + TCVT(row, f32, pto::RoundMode::CAST_NONE); + } + } + PIPE_BARRIER_VEC(); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + // Phase C (MTE3): store to convStates[ci, 0:K-1, c0:]. + for (int32_t i = 0; i < halo; ++i) { + IoTile row(lanes); + TASSIGN(row, i * ioTileBytes); + GlobalIoTensor dG(convStates + ((uint64_t)ci * stateLen + i) * dim + c0, {lanes}); + TSTORE(dG, row); + } + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); // store done -> UB free for the next strided task } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); // drain the final store before the kernel exits } + +} // namespace cc1d + +// Vector-only kernel (the cube/AIC pass gets empty bodies). One conv + writeback +// entry per (RS, dtype): templated on the compile-time ring size RS and tile width +// MAX_W, with the width K passed at runtime; the host launches the rs variant. +#if defined(__DAV_VEC__) +#define CONV_BODY(T, RS, MW) \ + cc1d::runConv((__gm__ T *)x, (__gm__ T *)wgt, (__gm__ T *)bia, (__gm__ T *)convStates, \ + (__gm__ int32_t *)qsl, (__gm__ int32_t *)cidx, (__gm__ uint8_t *)hinit, (__gm__ T *)y, \ + dim, batch, inputMode, seqLen, stateLen, width, col_w, blocksPerSeq, lchunks, activation, \ + hasBias, padSlot) +#define WB_BODY(T, RS, MW) \ + cc1d::runWriteback((__gm__ T *)x, (__gm__ T *)convStates, (__gm__ int32_t *)qsl, \ + (__gm__ int32_t *)cidx, (__gm__ uint8_t *)hinit, dim, batch, inputMode, seqLen, \ + stateLen, width, col_w, blocksPerSeq, padSlot) +#else // cube pass: empty bodies; void the params to silence unused warnings. +#define CONV_BODY(T, RS, MW) \ + (void)x, (void)wgt, (void)bia, (void)convStates, (void)qsl, (void)cidx, (void)hinit, (void)y, (void)dim, \ + (void)batch, (void)inputMode, (void)seqLen, (void)stateLen, (void)width, (void)col_w, (void)blocksPerSeq, \ + (void)lchunks, (void)activation, (void)hasBias, (void)padSlot +#define WB_BODY(T, RS, MW) \ + (void)x, (void)convStates, (void)qsl, (void)cidx, (void)hinit, (void)dim, (void)batch, (void)inputMode, \ + (void)seqLen, (void)stateLen, (void)width, (void)col_w, (void)blocksPerSeq, (void)padSlot +#endif + +#define CONV_PARAMS \ + GM_ADDR x, GM_ADDR wgt, GM_ADDR bia, GM_ADDR convStates, GM_ADDR qsl, GM_ADDR cidx, GM_ADDR hinit, GM_ADDR y, \ + uint32_t dim, uint32_t batch, uint32_t inputMode, uint32_t seqLen, uint32_t stateLen, uint32_t width, \ + uint32_t col_w, uint32_t blocksPerSeq, uint32_t lchunks, uint32_t activation, uint32_t hasBias, \ + int32_t padSlot +#define WB_PARAMS \ + GM_ADDR x, GM_ADDR convStates, GM_ADDR qsl, GM_ADDR cidx, GM_ADDR hinit, uint32_t dim, uint32_t batch, \ + uint32_t inputMode, uint32_t seqLen, uint32_t stateLen, uint32_t width, uint32_t col_w, uint32_t blocksPerSeq, \ + int32_t padSlot + +#define DEF_CONV(SUF, T, RS, MW) \ + extern "C" __global__ AICORE void causal_conv1d_##SUF(CONV_PARAMS) \ + { \ + CONV_BODY(T, RS, MW); \ + } +#define DEF_WB(SUF, T, RS, MW) \ + extern "C" __global__ AICORE void causal_conv1d_wb_##SUF(WB_PARAMS) \ + { \ + WB_BODY(T, RS, MW); \ + } + +// Ring sizes the kernel is compiled for -- must match FOR_EACH_RING_SIZE in +// op_host/causal_conv1d.cpp. Each row is (ringSize, maxTileWidth); a larger ring uses +// a smaller tile. A5 (dav-c310) has 256 KiB UB vs A2/A3's 192 KiB, so its tiles are +// ~4/3 wider -> fewer channel-tiles for large dim (e.g. dim=4096,K=4: 1 tile not 2). +// Both variants are checked against UB_BYTES_PER_CORE by the convChunk static_assert. +#ifdef __DAV_C310__ +#define FOR_EACH_RING_SIZE(DO) DO(2, 5120) DO(4, 4096) DO(8, 2048) DO(16, 1152) DO(32, 512) DO(64, 128) +#else +#define FOR_EACH_RING_SIZE(DO) DO(2, 4096) DO(4, 3072) DO(8, 1536) DO(16, 896) DO(32, 384) DO(64, 128) +#endif +#define DEFINE_ENTRIES(ringSize, maxTileWidth) \ + DEF_CONV(rs##ringSize##_half, half, ringSize, maxTileWidth) \ + DEF_CONV(rs##ringSize##_bf16, bfloat16_t, ringSize, maxTileWidth) \ + DEF_WB(rs##ringSize##_half, half, ringSize, maxTileWidth) \ + DEF_WB(rs##ringSize##_bf16, bfloat16_t, ringSize, maxTileWidth) +FOR_EACH_RING_SIZE(DEFINE_ENTRIES) +#undef DEFINE_ENTRIES +#undef FOR_EACH_RING_SIZE diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d.h b/csrc/causal_conv1d/op_kernel/causal_conv1d.h deleted file mode 100644 index 2689ed443..000000000 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d.h +++ /dev/null @@ -1,1021 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -/*! - * \file causal_conv1d.h - */ - -#ifndef CUSTOM_CAUSAL_CONV1D_H -#define CUSTOM_CAUSAL_CONV1D_H - -#include "kernel_operator.h" -#include "kernel_tiling/kernel_tiling.h" -#include "causal_conv1d_tiling_data.h" -#include "causal_conv1d_tiling_key.h" -#include "causal_conv1d_common.h" -#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 -#include "arch35/causal_conv1d_regbase.h" -#endif - -namespace NsCausalConv1d { - -using namespace AscendC; -using namespace NsCausalConv1dCommon; - -#define CAUSAL_CONV1D_TEMPLATE_ARGS typename T, uint32_t runModeKey, uint32_t widthKey, uint32_t fnPlanKey -#define CAUSAL_CONV1D_CLASS CausalConv1d - -enum SeqTaskWindowMode : int32_t { - SEQ_TASK_WINDOW_MODE_VARLEN = 0, - SEQ_TASK_WINDOW_MODE_BATCH = 1, - SEQ_TASK_WINDOW_MODE_DECODE2D = 2, -}; - -inline constexpr int32_t INIT_STATE_SYNCALL_NEED_SIZE = 8; -inline constexpr int32_t INIT_STATE_SYNCALL_MAX_BLOCKS = 64; - -struct SeqTaskWindow { - bool valid = false; - int32_t start = 0; - int32_t len = 0; -}; - -__aicore__ inline int32_t GetSeqTaskWindowMode(int32_t inputMode) -{ - if (inputMode == 0) { - return SEQ_TASK_WINDOW_MODE_VARLEN; - } - if (inputMode == 2) { - return SEQ_TASK_WINDOW_MODE_DECODE2D; - } - return SEQ_TASK_WINDOW_MODE_BATCH; -} - -__aicore__ inline SeqTaskWindow BuildSeqTaskWindowVarlen(int32_t startVal, int32_t endVal) -{ - SeqTaskWindow window; - window.start = startVal; - window.len = endVal - startVal; - window.valid = (window.len > 0); - return window; -} - -__aicore__ inline int32_t RetreatRingSlot(int32_t slot, int32_t delta) -{ - int32_t prev = slot - delta; - return (prev >= 0) ? prev : (prev + RING_SLOTS); -} - -__aicore__ inline SeqTaskWindow BuildSeqTaskWindowBatch(int32_t seq, int32_t seqLen) -{ - SeqTaskWindow window; - window.start = seq * seqLen; - window.len = seqLen; - window.valid = (window.len > 0); - return window; -} - -__aicore__ inline SeqTaskWindow BuildSeqTaskWindowDecode2D(int32_t seq) -{ - SeqTaskWindow window; - window.valid = true; - window.start = seq; - window.len = 1; - return window; -} - -__aicore__ inline constexpr int32_t DecodeWidthTplKey(uint32_t widthKey) -{ - switch (widthKey) { - case CAUSAL_CONV1D_TPL_WIDTH_2: - return 2; - case CAUSAL_CONV1D_TPL_WIDTH_3: - return 3; - case CAUSAL_CONV1D_TPL_WIDTH_4: - return 4; - default: - return 0; - } -} - -template -class CausalConv1d -{ -public: - __aicore__ inline CausalConv1d() = default; - - __aicore__ inline void ResetRuntimeState(const __gm__ CausalConv1dTilingData *tilingData) - { - tilingData_ = tilingData; - } - -protected: - static constexpr bool kIsUpdateMode = (runModeKey == CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE); - static constexpr int32_t kTemplateWidth = DecodeWidthTplKey(widthKey); - static constexpr bool kHasCompileTimeWidth = - (runModeKey == CAUSAL_CONV1D_TPL_RUN_MODE_FN) && (kTemplateWidth >= 2) && (kTemplateWidth <= MAX_WIDTH); - static constexpr FnExecutionPlan kFnExecutionPlan = static_cast(fnPlanKey); - - __aicore__ inline void InitSharedBuffersAndEvents(); - __aicore__ inline void LoadWeightAndBias(int32_t channelStart, int32_t baseDim); - __aicore__ inline void InitRing(int32_t cacheIdx, bool hasInit, int32_t stateTokenOffset, int32_t start, - int32_t len, int32_t channelStart, int32_t baseDim, int32_t dim); - __aicore__ inline void InitRingSeqSplit(int32_t seq, int32_t cacheIdx, bool hasInit, int32_t seqStart, - int32_t tileStart, int32_t tileLen, int32_t channelStart, int32_t baseDim, - int32_t dim); - __aicore__ inline void PrefetchInitStatesToWorkspace(int32_t channelStart, int32_t baseDimSize); - __aicore__ inline void RestoreFnLocalPartials(int32_t baseDim); - __aicore__ inline void ComputeFnRollingOutput(int32_t slotCurr, int32_t baseDim); - __aicore__ inline void AdvanceFnLocalPartials(int32_t slotCurr, int32_t baseDim); - __aicore__ inline void RunSeqFnRolling(int32_t start, int32_t len, int32_t channelStart, int32_t baseDim, - int32_t dim); - __aicore__ inline void RunSeq(int32_t start, int32_t len, int32_t channelStart, int32_t baseDim, int32_t dim); - __aicore__ inline void WriteBackState(int32_t cacheIdx, int32_t len, int32_t channelStart, int32_t baseDim, - int32_t dim); - __aicore__ inline void WriteBackStateSpec(int32_t cacheIdx, bool hasInit, int32_t stateTokenOffset, int32_t start, - int32_t len, int32_t channelStart, int32_t baseDim, int32_t dim); - __aicore__ inline void DrainTaskMte3(); - __aicore__ inline void AllocEvents(); - __aicore__ inline void ReleaseEvents(); - __aicore__ inline int32_t FindVarlenSeqByToken(int32_t tokenIdx) const; - __aicore__ inline bool ResolveExplicitTokenTileSeqRange(int32_t tokenTileId, int32_t &startSeq, - int32_t &endSeq) const; - __aicore__ inline bool ResolveSeqTaskWindow(int32_t seq, int32_t inputMode, int32_t seqLen, int32_t &start, - int32_t &len) const; - template - __aicore__ inline bool ResolveSeqTaskWindowByMode(int32_t seq, int32_t seqLen, int32_t &start, int32_t &len) const; - __aicore__ inline bool ResolveSeqCacheIndex(int32_t seq, bool hasCacheIndices, int32_t &cacheIdx) const; - __aicore__ inline bool ResolveSeqHasInit(int32_t seq, bool hasInitialStateMode) const; - __aicore__ inline void MaybeWriteBackSeqSplitTailChunk(int32_t chunkStart, int32_t chunkLen, int32_t seqStart, - int32_t seqLen, int32_t cacheIdx, int32_t channelStart, - int32_t baseDim, int32_t dim); - __aicore__ inline void ProcessDefault(); - template - __aicore__ inline void ProcessDefaultByWindowMode(); - __aicore__ inline void ProcessVarlenTokenTiled(); - __aicore__ inline void ProcessFnChunk(int32_t seq, int32_t cacheIdx, bool hasInit, int32_t seqStart, int32_t seqLen, - int32_t chunkStart, int32_t chunkLen, int32_t channelStart, int32_t baseDim, - int32_t dim); - __aicore__ inline const __gm__ CausalConv1dTilingData *GetTilingData() const; - __aicore__ inline bool HasActivation() const; - __aicore__ inline bool HasBias() const; - __aicore__ inline bool IsUpdateMode() const; - __aicore__ inline bool IsFnRollingFastPathEnabled() const; - __aicore__ inline bool HasExplicitFnTokenSeqRanges() const; - __aicore__ inline bool IsUpdateSpecDecodingEnabled() const; - -protected: - TPipe pipe; - TBuf inBuf; - TBuf outBuf; - TBuf calcBuf; - - TEventID weightBiasMte2ToVEvent_; - TEventID stateMte2ToVEvent_; - TEventID inputMte2ToVEvent_[RING_SLOTS]; - TEventID inputVToMte2Event_; - TEventID outMte3ToVEvent_[2]; - TEventID outVToMte3Event_[2]; - TEventID stateWritebackMte3ToVEvent_; - TEventID stateWritebackMte3ToMte2Event_; - TEventID stateShiftMte2ToMte3Event_; - TEventID stateShiftVToMte3Event_; - TEventID stateShiftMte3ToMte2Event_; - TEventID initSnapshotMte2ToMte3Event_; - TEventID initSnapshotMte3ToMte2Event_; - TEventID initSyncVToMte3Event_; - TEventID initSyncMte3ToVEvent_; - TEventID specWritebackMte2ToMte3Event_[2]; - TEventID specWritebackMte3ToMte2Event_[2]; - - GlobalTensor xGm; - GlobalTensor weightGm; - GlobalTensor biasGm; - GlobalTensor convStatesGm; - GlobalTensor queryStartLocGm; - GlobalTensor cacheIndicesGm; - GlobalTensor initialStateModeGm; - GlobalTensor numAcceptedTokensGm; - GlobalTensor yGm; - GlobalTensor initStateSyncGm_; - GlobalTensor initStateWorkspaceGm_; - - const __gm__ CausalConv1dTilingData *tilingData_{nullptr}; -}; - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::InitSharedBuffersAndEvents() -{ - pipe.InitBuffer(inBuf, RING_SLOTS * MAX_BLOCK_DIM * sizeof(T)); - pipe.InitBuffer(outBuf, 2 * MAX_BLOCK_DIM * sizeof(T)); - pipe.InitBuffer(calcBuf, (MAX_WIDTH + 4) * MAX_BLOCK_DIM * sizeof(float)); - AllocEvents(); -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::AllocEvents() -{ - weightBiasMte2ToVEvent_ = GetTPipePtr()->AllocEventID(); - stateMte2ToVEvent_ = GetTPipePtr()->AllocEventID(); - for (int32_t i = 0; i < RING_SLOTS; ++i) { - inputMte2ToVEvent_[i] = GetTPipePtr()->AllocEventID(); - } - inputVToMte2Event_ = GetTPipePtr()->AllocEventID(); - outMte3ToVEvent_[0] = GetTPipePtr()->AllocEventID(); - outMte3ToVEvent_[1] = GetTPipePtr()->AllocEventID(); - outVToMte3Event_[0] = GetTPipePtr()->AllocEventID(); - outVToMte3Event_[1] = GetTPipePtr()->AllocEventID(); - stateWritebackMte3ToVEvent_ = GetTPipePtr()->AllocEventID(); - stateWritebackMte3ToMte2Event_ = GetTPipePtr()->AllocEventID(); - stateShiftMte2ToMte3Event_ = GetTPipePtr()->AllocEventID(); - stateShiftVToMte3Event_ = GetTPipePtr()->AllocEventID(); - stateShiftMte3ToMte2Event_ = GetTPipePtr()->AllocEventID(); - initSnapshotMte2ToMte3Event_ = GetTPipePtr()->AllocEventID(); - initSnapshotMte3ToMte2Event_ = GetTPipePtr()->AllocEventID(); - initSyncVToMte3Event_ = GetTPipePtr()->AllocEventID(); - initSyncMte3ToVEvent_ = GetTPipePtr()->AllocEventID(); - specWritebackMte2ToMte3Event_[0] = GetTPipePtr()->AllocEventID(); - specWritebackMte2ToMte3Event_[1] = GetTPipePtr()->AllocEventID(); - specWritebackMte3ToMte2Event_[0] = GetTPipePtr()->AllocEventID(); - specWritebackMte3ToMte2Event_[1] = GetTPipePtr()->AllocEventID(); -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::ReleaseEvents() -{ - GetTPipePtr()->ReleaseEventID(weightBiasMte2ToVEvent_); - GetTPipePtr()->ReleaseEventID(stateMte2ToVEvent_); - for (int32_t i = 0; i < RING_SLOTS; ++i) { - GetTPipePtr()->ReleaseEventID(inputMte2ToVEvent_[i]); - } - GetTPipePtr()->ReleaseEventID(inputVToMte2Event_); - GetTPipePtr()->ReleaseEventID(outMte3ToVEvent_[0]); - GetTPipePtr()->ReleaseEventID(outMte3ToVEvent_[1]); - GetTPipePtr()->ReleaseEventID(outVToMte3Event_[0]); - GetTPipePtr()->ReleaseEventID(outVToMte3Event_[1]); - GetTPipePtr()->ReleaseEventID(stateWritebackMte3ToVEvent_); - GetTPipePtr()->ReleaseEventID(stateWritebackMte3ToMte2Event_); - GetTPipePtr()->ReleaseEventID(stateShiftMte2ToMte3Event_); - GetTPipePtr()->ReleaseEventID(stateShiftVToMte3Event_); - GetTPipePtr()->ReleaseEventID(stateShiftMte3ToMte2Event_); - GetTPipePtr()->ReleaseEventID(initSnapshotMte2ToMte3Event_); - GetTPipePtr()->ReleaseEventID(initSnapshotMte3ToMte2Event_); - GetTPipePtr()->ReleaseEventID(initSyncVToMte3Event_); - GetTPipePtr()->ReleaseEventID(initSyncMte3ToVEvent_); - GetTPipePtr()->ReleaseEventID(specWritebackMte2ToMte3Event_[0]); - GetTPipePtr()->ReleaseEventID(specWritebackMte2ToMte3Event_[1]); - GetTPipePtr()->ReleaseEventID(specWritebackMte3ToMte2Event_[0]); - GetTPipePtr()->ReleaseEventID(specWritebackMte3ToMte2Event_[1]); -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::LoadWeightAndBias(int32_t channelStart, int32_t baseDim) -{ - const int32_t dim = tilingData_->dim; - const int32_t width = static_cast(tilingData_->width); - const int32_t jStart = MAX_WIDTH - width; - const bool hasBias = HasBias(); - auto cl = CalcBufLayout::FromCalcBuf(calcBuf); - LocalTensor &weightF = cl.weightF; - LocalTensor &biasF = cl.biasF; - LocalTensor weightT; - LocalTensor biasT; - - if constexpr (!std::is_same::value) { - weightT = weightF.ReinterpretCast(); - biasT = biasF.ReinterpretCast(); - } - - for (int32_t j = 0; j < jStart; ++j) { - Duplicate(weightF[j * MAX_BLOCK_DIM], 0.0f, baseDim); - } - - for (int32_t j = 0; j < width; ++j) { - const int32_t jDst = jStart + j; - const int64_t weightOffset = static_cast(j) * dim + channelStart; - - if constexpr (std::is_same::value) { - DataCopy(weightF[jDst * MAX_BLOCK_DIM], weightGm[weightOffset], baseDim); - } else { - DataCopy(weightT[jDst * MAX_BLOCK_DIM * 2 + MAX_BLOCK_DIM], weightGm[weightOffset], baseDim); - } - } - - if (hasBias) { - if constexpr (std::is_same::value) { - DataCopy(biasF, biasGm[channelStart], baseDim); - } else { - DataCopy(biasT[MAX_BLOCK_DIM], biasGm[channelStart], baseDim); - } - } - - SetFlag(weightBiasMte2ToVEvent_); - WaitFlag(weightBiasMte2ToVEvent_); - - if constexpr (!std::is_same::value) { - for (int32_t j = 0; j < width; ++j) { - const int32_t jDst = jStart + j; - Cast(weightF[jDst * MAX_BLOCK_DIM], weightT[jDst * MAX_BLOCK_DIM * 2 + MAX_BLOCK_DIM], RoundMode::CAST_NONE, - baseDim); - } - if (hasBias) { - Cast(biasF, biasT[MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - } - PipeBarrier(); - } - - if (!hasBias) { - Duplicate(biasF, 0.0f, baseDim); - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::InitRing(int32_t cacheIdx, bool hasInit, int32_t stateTokenOffset, - int32_t start, int32_t len, int32_t channelStart, int32_t baseDim, - int32_t dim) -{ - const int32_t stateLen = tilingData_->stateLen; - const int32_t width = static_cast(tilingData_->width); - const int32_t ringStart = MAX_WIDTH - width; - LocalTensor ring = inBuf.Get(); - - for (int32_t i = 0; i < ringStart; ++i) { - Duplicate(ring[i * MAX_BLOCK_DIM], static_cast(0), baseDim); - } - if (ringStart > 0) { - PipeBarrier(); - } - - if (hasInit) { - for (int32_t i = 0; i < (width - 1); ++i) { - const int32_t pos = stateTokenOffset + i; - const int64_t stateOffset = - static_cast(cacheIdx) * stateLen * dim + static_cast(pos) * dim + channelStart; - DataCopy(ring[(ringStart + i) * MAX_BLOCK_DIM], convStatesGm[stateOffset], baseDim); - } - SetFlag(stateMte2ToVEvent_); - WaitFlag(stateMte2ToVEvent_); - } else { - for (int32_t i = 0; i < (width - 1); ++i) { - Duplicate(ring[(ringStart + i) * MAX_BLOCK_DIM], static_cast(0), baseDim); - } - PipeBarrier(); - } - - if (len > 0) { - const int32_t slot0 = SlotCurr(0); - const int64_t xOffset = static_cast(start) * dim + channelStart; - DataCopy(ring[slot0 * MAX_BLOCK_DIM], xGm[xOffset], baseDim); - SetFlag(inputMte2ToVEvent_[slot0]); - } - - if (len > 1) { - SetFlag(inputVToMte2Event_); - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeq(int32_t start, int32_t len, int32_t channelStart, int32_t baseDim, - int32_t dim) -{ - if (IsFnRollingFastPathEnabled()) { - RunSeqFnRolling(start, len, channelStart, baseDim, dim); - return; - } - - const int32_t width = static_cast(tilingData_->width); - const int32_t jStart = MAX_WIDTH - width; - auto cl = CalcBufLayout::FromCalcBuf(calcBuf); - LocalTensor &weightF = cl.weightF; - LocalTensor &biasF = cl.biasF; - LocalTensor &accF = cl.accF; - LocalTensor &tmpF = cl.tmpF; - LocalTensor ring = inBuf.Get(); - LocalTensor outT = outBuf.Get(); - const bool hasBias = HasBias(); - const bool hasActivation = HasActivation(); - for (int32_t t = 0; t < len; ++t) { - const int32_t slotCurr = SlotCurr(t); - - WaitFlag(inputMte2ToVEvent_[slotCurr]); - - if (t + 1 < len) { - const int32_t slotNext = SlotPrefetch(t); - const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; - WaitFlag(inputVToMte2Event_); - DataCopy(ring[slotNext * MAX_BLOCK_DIM], xGm[xOffsetNext], baseDim); - SetFlag(inputMte2ToVEvent_[slotNext]); - } - - bool accInitialized = false; - if (hasBias) { - Adds(accF, biasF, 0.0f, baseDim); - PipeBarrier(); - accInitialized = true; - } - - for (int32_t j = jStart; j < MAX_WIDTH; ++j) { - const int32_t tap = (MAX_WIDTH - 1) - j; - const int32_t slot = (tap == 0) ? slotCurr : SlotHist(t, tap); - Cast(tmpF, ring[slot * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - if (!accInitialized) { - Mul(accF, tmpF, weightF[j * MAX_BLOCK_DIM], baseDim); - accInitialized = true; - } else { - MulAddDst(accF, tmpF, weightF[j * MAX_BLOCK_DIM], baseDim); - } - } - - PipeBarrier(); - - if (hasActivation) { - Silu(tmpF, accF, baseDim); - } - - const int32_t outSlot = t & 1; - LocalTensor outSlotT = outT[outSlot * MAX_BLOCK_DIM]; - if (t >= 2) { - WaitFlag(outMte3ToVEvent_[outSlot]); - } - - if constexpr (IsSameType::value) { - if (hasActivation) { - DataCopy(outSlotT, tmpF, baseDim); - } else { - DataCopy(outSlotT, accF, baseDim); - } - } else { - if (hasActivation) { - Cast(outSlotT, tmpF, RoundMode::CAST_RINT, baseDim); - } else { - Cast(outSlotT, accF, RoundMode::CAST_RINT, baseDim); - } - } - - SetFlag(outVToMte3Event_[outSlot]); - - const int64_t outOffset = static_cast(start + t) * dim + channelStart; - WaitFlag(outVToMte3Event_[outSlot]); - DataCopy(yGm[outOffset], outSlotT, baseDim); - if (t + 2 < len) { - SetFlag(outMte3ToVEvent_[outSlot]); - } - - if (t + 2 < len) { - SetFlag(inputVToMte2Event_); - } - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::RestoreFnLocalPartials(int32_t baseDim) -{ - if constexpr (!kHasCompileTimeWidth) { - return; - } - - auto cl = CalcBufLayout::FromCalcBuf(calcBuf); - LocalTensor &weightF = cl.weightF; - LocalTensor &state2F = cl.biasF; - LocalTensor &state1F = cl.accF; - LocalTensor &state0F = cl.tmpF; - LocalTensor &currF = cl.currF; - LocalTensor ring = inBuf.Get(); - constexpr int32_t ringStart = MAX_WIDTH - kTemplateWidth; - constexpr int32_t w0Idx = MAX_WIDTH - kTemplateWidth; - - if constexpr (kTemplateWidth == 2) { - Duplicate(state2F, 0.0f, baseDim); - Duplicate(state1F, 0.0f, baseDim); - PipeBarrier(); - - Cast(currF, ring[ringStart * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - } else if constexpr (kTemplateWidth == 3) { - Duplicate(state2F, 0.0f, baseDim); - PipeBarrier(); - - Cast(currF, ring[ringStart * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - - Cast(currF, ring[(ringStart + 1) * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - Mul(state1F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - MulAddDst(state0F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - } else if constexpr (kTemplateWidth == 4) { - Cast(currF, ring[ringStart * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - - Cast(currF, ring[(ringStart + 1) * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - Mul(state1F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - MulAddDst(state0F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - - Cast(currF, ring[(ringStart + 2) * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - Mul(state2F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - MulAddDst(state1F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - MulAddDst(state0F, currF, weightF[(w0Idx + 2) * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::ComputeFnRollingOutput(int32_t slotCurr, int32_t baseDim) -{ - if constexpr (!kHasCompileTimeWidth) { - return; - } - - auto cl = CalcBufLayout::FromCalcBuf(calcBuf); - LocalTensor &weightF = cl.weightF; - LocalTensor &state0F = cl.tmpF; - LocalTensor &currF = cl.currF; - LocalTensor ring = inBuf.Get(); - -#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 - const bool hasActivation = HasActivation(); - if (hasActivation) { - ComputeFnRollingOutputRegbase(ring[slotCurr * MAX_BLOCK_DIM], currF, state0F, - weightF[3 * MAX_BLOCK_DIM], baseDim); - } else { - ComputeFnRollingOutputRegbase(ring[slotCurr * MAX_BLOCK_DIM], currF, state0F, - weightF[3 * MAX_BLOCK_DIM], baseDim); - } -#else - Cast(currF, ring[slotCurr * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - MulAddDst(state0F, currF, weightF[3 * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - - const bool hasActivation = HasActivation(); - if (hasActivation) { - PipeBarrier(); - Silu(currF, state0F, baseDim); - } -#endif -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::AdvanceFnLocalPartials(int32_t slotCurr, int32_t baseDim) -{ - if constexpr (!kHasCompileTimeWidth) { - return; - } - - auto cl = CalcBufLayout::FromCalcBuf(calcBuf); - LocalTensor &weightF = cl.weightF; - LocalTensor &state2F = cl.biasF; - LocalTensor &state1F = cl.accF; - LocalTensor &state0F = cl.tmpF; - LocalTensor &currF = cl.currF; - LocalTensor ring = inBuf.Get(); - constexpr int32_t w0Idx = MAX_WIDTH - kTemplateWidth; - -#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 - AdvanceFnLocalPartialsRegbase(ring[slotCurr * MAX_BLOCK_DIM], weightF[w0Idx * MAX_BLOCK_DIM], - state0F, state1F, state2F, baseDim, MAX_BLOCK_DIM); -#else - Cast(currF, ring[slotCurr * MAX_BLOCK_DIM], RoundMode::CAST_NONE, baseDim); - PipeBarrier(); - - if constexpr (kTemplateWidth == 2) { - Mul(state0F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - } else if constexpr (kTemplateWidth == 3) { - Mul(state0F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - Add(state0F, state0F, state1F, baseDim); - PipeBarrier(); - - Mul(state1F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - } else if constexpr (kTemplateWidth == 4) { - Mul(state0F, currF, weightF[(w0Idx + 2) * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - Add(state0F, state0F, state1F, baseDim); - PipeBarrier(); - - Mul(state1F, currF, weightF[(w0Idx + 1) * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - Add(state1F, state1F, state2F, baseDim); - PipeBarrier(); - - Mul(state2F, currF, weightF[w0Idx * MAX_BLOCK_DIM], baseDim); - PipeBarrier(); - } -#endif -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::RunSeqFnRolling(int32_t start, int32_t len, int32_t channelStart, - int32_t baseDim, int32_t dim) -{ - if constexpr (!kHasCompileTimeWidth) { - return; - } - - auto cl = CalcBufLayout::FromCalcBuf(calcBuf); - LocalTensor &state0F = cl.tmpF; - LocalTensor &currF = cl.currF; - LocalTensor ring = inBuf.Get(); - LocalTensor outT = outBuf.Get(); - const bool hasActivation = HasActivation(); - RestoreFnLocalPartials(baseDim); - - for (int32_t t = 0; t < len; ++t) { - const int32_t slotCurr = SlotCurr(t); - - WaitFlag(inputMte2ToVEvent_[slotCurr]); - - if (t + 1 < len) { - const int32_t slotNext = SlotPrefetch(t); - const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; - WaitFlag(inputVToMte2Event_); - DataCopy(ring[slotNext * MAX_BLOCK_DIM], xGm[xOffsetNext], baseDim); - SetFlag(inputMte2ToVEvent_[slotNext]); - } - - ComputeFnRollingOutput(slotCurr, baseDim); - - const int32_t outSlot = t & 1; - LocalTensor outSlotT = outT[outSlot * MAX_BLOCK_DIM]; - if (t >= 2) { - WaitFlag(outMte3ToVEvent_[outSlot]); - } - - if constexpr (IsSameType::value) { - if (hasActivation) { - DataCopy(outSlotT, currF, baseDim); - } else { - DataCopy(outSlotT, state0F, baseDim); - } - } else { - if (hasActivation) { - Cast(outSlotT, currF, RoundMode::CAST_RINT, baseDim); - } else { - Cast(outSlotT, state0F, RoundMode::CAST_RINT, baseDim); - } - } - - AdvanceFnLocalPartials(slotCurr, baseDim); - - SetFlag(outVToMte3Event_[outSlot]); - - const int64_t outOffset = static_cast(start + t) * dim + channelStart; - WaitFlag(outVToMte3Event_[outSlot]); - DataCopy(yGm[outOffset], outSlotT, baseDim); - if (t + 2 < len) { - SetFlag(outMte3ToVEvent_[outSlot]); - } - - if (t + 2 < len) { - SetFlag(inputVToMte2Event_); - } - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::DrainTaskMte3() -{ - SetFlag(stateWritebackMte3ToVEvent_); - WaitFlag(stateWritebackMte3ToVEvent_); - SetFlag(stateWritebackMte3ToMte2Event_); - WaitFlag(stateWritebackMte3ToMte2Event_); -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::WriteBackState(int32_t cacheIdx, int32_t len, int32_t channelStart, - int32_t baseDim, int32_t dim) -{ - const int32_t stateLen = tilingData_->stateLen; - const int32_t width = static_cast(tilingData_->width); - if (len <= 0) { - return; - } - - const int32_t lastT = len - 1; - LocalTensor ring = inBuf.Get(); - const int32_t lastSlot = SlotCurr(lastT); - const int64_t stateBaseOffset = static_cast(cacheIdx) * stateLen * dim + channelStart; - - for (int32_t pos = 0; pos < (width - 1); ++pos) { - const int32_t tap = (width - 2) - pos; - const int32_t slot = RetreatRingSlot(lastSlot, tap); - const int64_t stateOffset = stateBaseOffset + static_cast(pos) * dim; - DataCopy(convStatesGm[stateOffset], ring[slot * MAX_BLOCK_DIM], baseDim); - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::WriteBackStateSpec(int32_t cacheIdx, bool hasInit, int32_t stateTokenOffset, - int32_t start, int32_t len, int32_t channelStart, - int32_t baseDim, int32_t dim) -{ - const int32_t width = static_cast(tilingData_->width); - const int32_t stateLen = tilingData_->stateLen; - if (len <= 0) { - return; - } - - if (width != 4) { - WriteBackState(cacheIdx, len, channelStart, baseDim, dim); - return; - } - - constexpr int32_t keep = MAX_WIDTH - 2; - const int32_t reqStateLen = keep + len; - if (reqStateLen > stateLen) { - WriteBackState(cacheIdx, len, channelStart, baseDim, dim); - return; - } - - LocalTensor ring = inBuf.Get(); - LocalTensor buf0 = ring[0 * MAX_BLOCK_DIM]; - LocalTensor buf1 = ring[1 * MAX_BLOCK_DIM]; - - if (hasInit) { - const int32_t srcPos0 = stateTokenOffset + 1; - const int32_t srcPos1 = stateTokenOffset + 2; - const int64_t srcOffset0 = - static_cast(cacheIdx) * stateLen * dim + static_cast(srcPos0) * dim + channelStart; - const int64_t srcOffset1 = - static_cast(cacheIdx) * stateLen * dim + static_cast(srcPos1) * dim + channelStart; - DataCopy(buf0, convStatesGm[srcOffset0], baseDim); - DataCopy(buf1, convStatesGm[srcOffset1], baseDim); - SetFlag(stateShiftMte2ToMte3Event_); - WaitFlag(stateShiftMte2ToMte3Event_); - const int64_t dstOffset0 = - static_cast(cacheIdx) * stateLen * dim + static_cast(0) * dim + channelStart; - const int64_t dstOffset1 = - static_cast(cacheIdx) * stateLen * dim + static_cast(1) * dim + channelStart; - DataCopy(convStatesGm[dstOffset0], buf0, baseDim); - DataCopy(convStatesGm[dstOffset1], buf1, baseDim); - SetFlag(stateShiftMte3ToMte2Event_); - WaitFlag(stateShiftMte3ToMte2Event_); - } else { - Duplicate(buf0, static_cast(0), baseDim); - SetFlag(stateShiftVToMte3Event_); - WaitFlag(stateShiftVToMte3Event_); - const int64_t dstOffset0 = - static_cast(cacheIdx) * stateLen * dim + static_cast(0) * dim + channelStart; - const int64_t dstOffset1 = - static_cast(cacheIdx) * stateLen * dim + static_cast(1) * dim + channelStart; - DataCopy(convStatesGm[dstOffset0], buf0, baseDim); - DataCopy(convStatesGm[dstOffset1], buf0, baseDim); - SetFlag(stateShiftMte3ToMte2Event_); - WaitFlag(stateShiftMte3ToMte2Event_); - } - - const int64_t xOffset0 = static_cast(start) * dim + channelStart; - DataCopy(buf0, xGm[xOffset0], baseDim); - SetFlag(specWritebackMte2ToMte3Event_[0]); - - for (int32_t t = 0; t < len; ++t) { - const int32_t curr = t & 1; - const int32_t next = curr ^ 1; - LocalTensor currBuf = (curr == 0) ? buf0 : buf1; - LocalTensor nextBuf = (next == 0) ? buf0 : buf1; - - WaitFlag(specWritebackMte2ToMte3Event_[curr]); - - if (t + 1 < len) { - const int64_t xOffsetNext = static_cast(start + t + 1) * dim + channelStart; - if (t > 0) { - WaitFlag(specWritebackMte3ToMte2Event_[next]); - } - DataCopy(nextBuf, xGm[xOffsetNext], baseDim); - SetFlag(specWritebackMte2ToMte3Event_[next]); - } - - const int64_t dstOffset = - static_cast(cacheIdx) * stateLen * dim + static_cast(keep + t) * dim + channelStart; - DataCopy(convStatesGm[dstOffset], currBuf, baseDim); - SetFlag(specWritebackMte3ToMte2Event_[curr]); - } - - WaitFlag(specWritebackMte3ToMte2Event_[0]); - if (len > 1) { - WaitFlag(specWritebackMte3ToMte2Event_[1]); - } -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqTaskWindow(int32_t seq, int32_t inputMode, int32_t seqLen, - int32_t &start, int32_t &len) const -{ - switch (GetSeqTaskWindowMode(inputMode)) { - case SEQ_TASK_WINDOW_MODE_VARLEN: - return ResolveSeqTaskWindowByMode(seq, seqLen, start, len); - case SEQ_TASK_WINDOW_MODE_DECODE2D: - return ResolveSeqTaskWindowByMode(seq, seqLen, start, len); - default: - return ResolveSeqTaskWindowByMode(seq, seqLen, start, len); - } -} - -template -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqTaskWindowByMode(int32_t seq, int32_t seqLen, int32_t &start, - int32_t &len) const -{ - SeqTaskWindow window; - if constexpr (kWindowMode == SEQ_TASK_WINDOW_MODE_VARLEN) { - const int32_t startVal = queryStartLocGm.GetValue(seq); - const int32_t endVal = queryStartLocGm.GetValue(seq + 1); - window = BuildSeqTaskWindowVarlen(startVal, endVal); - } else if constexpr (kWindowMode == SEQ_TASK_WINDOW_MODE_DECODE2D) { - window = BuildSeqTaskWindowDecode2D(seq); - } else { - window = BuildSeqTaskWindowBatch(seq, seqLen); - } - - if (!window.valid) { - return false; - } - start = window.start; - len = window.len; - return true; -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqCacheIndex(int32_t seq, bool hasCacheIndices, - int32_t &cacheIdx) const -{ - cacheIdx = seq; - if (!hasCacheIndices) { - return true; - } - - const int64_t cacheIdx64 = cacheIndicesGm.GetValue(seq); - if (cacheIdx64 == tilingData_->padSlotId) { - return false; - } - cacheIdx = static_cast(cacheIdx64); - return true; -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveSeqHasInit(int32_t seq, bool hasInitialStateMode) const -{ - return hasInitialStateMode ? (initialStateModeGm.GetValue(seq) != 0) : false; -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessDefault() -{ - switch (GetSeqTaskWindowMode(tilingData_->inputMode)) { - case SEQ_TASK_WINDOW_MODE_VARLEN: - ProcessDefaultByWindowMode(); - return; - case SEQ_TASK_WINDOW_MODE_DECODE2D: - ProcessDefaultByWindowMode(); - return; - default: - ProcessDefaultByWindowMode(); - return; - } -} - -template -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessDefaultByWindowMode() -{ - const int32_t dim = tilingData_->dim; - const int32_t batch = tilingData_->batch; - const int32_t seqLen = tilingData_->seqLen; - const int32_t baseDim = static_cast(tilingData_->baseDim); - const int32_t baseDimCnt = static_cast(tilingData_->baseDimCnt); - const int32_t width = static_cast(tilingData_->width); - const bool hasCacheIndices = (tilingData_->hasCacheIndices != 0); - const bool hasInit = true; - const bool isSpecDecodingGlobal = IsUpdateSpecDecodingEnabled(); - - const uint32_t blockIdx = GetBlockIdx(); - const uint32_t blockNum = GetBlockNum(); - - if (baseDim <= 0 || baseDimCnt <= 0 || baseDim > MAX_BLOCK_DIM || width < 2 || width > MAX_WIDTH) { - ReleaseEvents(); - return; - } - - const int64_t gridSize = static_cast(batch) * baseDimCnt; - for (int64_t task = static_cast(blockIdx); task < gridSize; task += static_cast(blockNum)) { - const int32_t seq = static_cast(task / baseDimCnt); - const int32_t baseDimIdx = static_cast(task % baseDimCnt); - const int32_t channelStart = baseDimIdx * baseDim; - if (channelStart >= dim) { - continue; - } - const int32_t curBaseDim = (channelStart + baseDim <= dim) ? baseDim : (dim - channelStart); - - int32_t start = 0; - int32_t len = 0; - if (!ResolveSeqTaskWindowByMode(seq, seqLen, start, len)) { - continue; - } - - int32_t cacheIdx = 0; - if (!ResolveSeqCacheIndex(seq, hasCacheIndices, cacheIdx)) { - continue; - } - - LoadWeightAndBias(channelStart, curBaseDim); - - if (isSpecDecodingGlobal) { - int32_t accepted = static_cast(numAcceptedTokensGm.GetValue(seq)); - int32_t stateTokenOffset = accepted - 1; - const int32_t maxOffset = static_cast(tilingData_->stateLen - (width - 1)); - if (stateTokenOffset < 0) { - stateTokenOffset = 0; - } else if (stateTokenOffset > maxOffset) { - stateTokenOffset = maxOffset; - } - - InitRing(cacheIdx, hasInit, stateTokenOffset, start, len, channelStart, curBaseDim, dim); - RunSeq(start, len, channelStart, curBaseDim, dim); - DrainTaskMte3(); - WriteBackStateSpec(cacheIdx, hasInit, stateTokenOffset, start, len, channelStart, curBaseDim, dim); - } else { - InitRing(cacheIdx, hasInit, 0, start, len, channelStart, curBaseDim, dim); - RunSeq(start, len, channelStart, curBaseDim, dim); - WriteBackState(cacheIdx, len, channelStart, curBaseDim, dim); - } - - DrainTaskMte3(); - } -} - -template -__aicore__ inline const __gm__ CausalConv1dTilingData *CAUSAL_CONV1D_CLASS::GetTilingData() const -{ - return tilingData_; -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::HasActivation() const -{ - return (tilingData_ != nullptr) && (tilingData_->activationMode != 0); -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::HasBias() const -{ - return (tilingData_ != nullptr) && (tilingData_->hasBias != 0); -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::IsUpdateMode() const -{ - return kIsUpdateMode; -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::IsFnRollingFastPathEnabled() const -{ - return !kIsUpdateMode && (tilingData_ != nullptr) && (kFnExecutionPlan != FN_EXECUTION_PLAN_INVALID) && - (tilingData_->hasNumAcceptedTokens == 0) && !HasBias(); -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::HasExplicitFnTokenSeqRanges() const -{ - return !kIsUpdateMode && (tilingData_ != nullptr) && (tilingData_->inputMode == 0) && - (tilingData_->hasExplicitTokenSeqRanges != 0) && - (tilingData_->explicitTokenSeqRangeCount >= tilingData_->tokenBlockCnt); -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::IsUpdateSpecDecodingEnabled() const -{ - return kIsUpdateMode && (tilingData_->hasNumAcceptedTokens != 0) && (tilingData_->width == 4); -} - -#include "causal_conv1d_fn_tasks.h" - -#undef CAUSAL_CONV1D_CLASS -#undef CAUSAL_CONV1D_TEMPLATE_ARGS - -} // namespace NsCausalConv1d -#endif // CUSTOM_CAUSAL_CONV1D_H diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_common.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_common.h deleted file mode 100644 index 5c94514ff..000000000 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d_common.h +++ /dev/null @@ -1,66 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -/*! - * \file causal_conv1d_common.h - */ - -#ifndef CUSTOM_CAUSAL_CONV1D_COMMON_H -#define CUSTOM_CAUSAL_CONV1D_COMMON_H - -#include "kernel_operator.h" - -namespace NsCausalConv1dCommon { - -constexpr int32_t MAX_WIDTH = 4; -constexpr int32_t MAX_BLOCK_DIM = 4096; -constexpr int32_t RING_SLOTS = 5; - -__aicore__ inline int32_t SlotCurr(int32_t t) -{ - return (t + 3) % RING_SLOTS; -} - -__aicore__ inline int32_t SlotHist(int32_t t, int32_t i) -{ - return (t + 3 - i) % RING_SLOTS; -} - -__aicore__ inline int32_t SlotPrefetch(int32_t t) -{ - return (t + 4) % RING_SLOTS; -} - -struct CalcBufLayout { - AscendC::LocalTensor weightF; - AscendC::LocalTensor biasF; - AscendC::LocalTensor accF; - AscendC::LocalTensor tmpF; - AscendC::LocalTensor currF; - - __aicore__ inline CalcBufLayout() = default; - - __aicore__ static inline CalcBufLayout FromCalcBuf(AscendC::TBuf &calcBuf) - { - CalcBufLayout layout; - AscendC::LocalTensor calc = calcBuf.template Get(); - layout.weightF = calc; - layout.biasF = calc[MAX_WIDTH * MAX_BLOCK_DIM]; - layout.accF = layout.biasF[MAX_BLOCK_DIM]; - layout.tmpF = layout.accF[MAX_BLOCK_DIM]; - layout.currF = layout.tmpF[MAX_BLOCK_DIM]; - return layout; - } -}; - -} // namespace NsCausalConv1dCommon - -#endif // CUSTOM_CAUSAL_CONV1D_COMMON_H diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_fn.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_fn.h deleted file mode 100644 index e484639a8..000000000 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d_fn.h +++ /dev/null @@ -1,72 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CUSTOM_CAUSAL_CONV1D_FN_H -#define CUSTOM_CAUSAL_CONV1D_FN_H - -#include "causal_conv1d.h" - -namespace NsCausalConv1d { - -template -class CausalConv1dFn : public CausalConv1d -{ -public: - __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, GM_ADDR queryStartLoc, - GM_ADDR cacheIndices, GM_ADDR initialStateMode, GM_ADDR numAcceptedTokens, GM_ADDR y, - GM_ADDR workspace, const __gm__ CausalConv1dTilingData *tilingData) - { - (void)numAcceptedTokens; - this->ResetRuntimeState(tilingData); - this->xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(x)); - this->weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(weight)); - this->biasGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(bias)); - this->convStatesGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(convStates)); - this->queryStartLocGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(queryStartLoc)); - this->cacheIndicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(cacheIndices)); - this->initialStateModeGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(initialStateMode)); - this->yGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(y)); - if (tilingData->hasInitStateWorkspace != 0) { - const uint64_t syncElems = static_cast(GetBlockNum()) * INIT_STATE_SYNCALL_NEED_SIZE; - const uint64_t syncBytes = syncElems * sizeof(int32_t); - const uint64_t workspaceElems = static_cast(tilingData->batch) * - static_cast(tilingData->width - 1) * - static_cast(tilingData->dim); - this->initStateSyncGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(workspace), syncElems); - auto *workspaceBytes = reinterpret_cast<__gm__ uint8_t *>(workspace); - this->initStateWorkspaceGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(workspaceBytes + syncBytes), - workspaceElems); - } - this->InitSharedBuffersAndEvents(); - } - - __aicore__ inline void Process() - { - this->ProcessVarlenTokenTiled(); - this->ReleaseEvents(); - } -}; - -template -__aicore__ inline void RunCausalConv1dFn(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, - GM_ADDR queryStartLoc, GM_ADDR cacheIndices, GM_ADDR initialStateMode, - GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, - const __gm__ CausalConv1dTilingData *tilingData) -{ - CausalConv1dFn op; - op.Init(x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, workspace, - tilingData); - op.Process(); -} - -} // namespace NsCausalConv1d - -#endif // CUSTOM_CAUSAL_CONV1D_FN_H diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h deleted file mode 100644 index d467dc478..000000000 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d_fn_tasks.h +++ /dev/null @@ -1,306 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CUSTOM_CAUSAL_CONV1D_FN_TASKS_H -#define CUSTOM_CAUSAL_CONV1D_FN_TASKS_H - -struct FnDirectBlockTask { - bool valid = false; - int32_t tokenTileId = 0; - int32_t baseDimIdx = 0; - int32_t tokenStart = 0; - int32_t tokenEnd = 0; - int32_t channelStart = 0; - int32_t baseDimSize = 0; -}; - -__aicore__ inline FnDirectBlockTask ResolveFnDirectBlockTask(int32_t blockIdx, int32_t tokenBlockCnt, - int32_t tokenBlockSize, int32_t cuSeqlen, - int32_t baseDimCnt, int32_t baseDim, int32_t dim) -{ - FnDirectBlockTask task; - if (blockIdx < 0 || tokenBlockCnt <= 0 || tokenBlockSize <= 0 || cuSeqlen <= 0 || baseDimCnt <= 0 || baseDim <= 0 || - dim <= 0) { - return task; - } - - const int64_t phase1Grid = static_cast(tokenBlockCnt) * baseDimCnt; - if (phase1Grid <= 0 || static_cast(blockIdx) >= phase1Grid) { - return task; - } - - task.tokenTileId = blockIdx / baseDimCnt; - task.baseDimIdx = blockIdx % baseDimCnt; - task.channelStart = task.baseDimIdx * baseDim; - if (task.channelStart >= dim) { - return task; - } - - task.baseDimSize = (task.channelStart + baseDim <= dim) ? baseDim : (dim - task.channelStart); - task.tokenStart = task.tokenTileId * tokenBlockSize; - if (task.tokenStart >= cuSeqlen) { - return task; - } - - const int32_t tokenEndRaw = task.tokenStart + tokenBlockSize; - task.tokenEnd = (tokenEndRaw <= cuSeqlen) ? tokenEndRaw : cuSeqlen; - if (task.baseDimSize <= 0 || task.tokenEnd <= task.tokenStart) { - return {}; - } - - task.valid = true; - return task; -} - -__aicore__ inline bool IsFnInitStateSnapshotOwnerBlock(const FnDirectBlockTask &task) -{ - return task.valid && task.tokenTileId == 0; -} - -template -__aicore__ inline int32_t CAUSAL_CONV1D_CLASS::FindVarlenSeqByToken(int32_t tokenIdx) const -{ - int32_t left = 0; - int32_t right = static_cast(tilingData_->batch); - while (left < right) { - const int32_t mid = left + ((right - left) >> 1); - const int32_t endVal = static_cast(queryStartLocGm.GetValue(mid + 1)); - if (tokenIdx < endVal) { - right = mid; - } else { - left = mid + 1; - } - } - return left; -} - -template -__aicore__ inline bool CAUSAL_CONV1D_CLASS::ResolveExplicitTokenTileSeqRange(int32_t tokenTileId, int32_t &startSeq, - int32_t &endSeq) const -{ - if (!HasExplicitFnTokenSeqRanges() || tokenTileId < 0 || tokenTileId >= tilingData_->explicitTokenSeqRangeCount) { - return false; - } - startSeq = static_cast(tilingData_->tokenTileStartSeq[tokenTileId]); - endSeq = static_cast(tilingData_->tokenTileEndSeq[tokenTileId]); - return (startSeq >= 0) && (endSeq >= startSeq); -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::InitRingSeqSplit(int32_t seq, int32_t cacheIdx, bool hasInit, - int32_t seqStart, int32_t tileStart, int32_t tileLen, - int32_t channelStart, int32_t baseDim, int32_t dim) -{ - const int32_t stateLen = tilingData_->stateLen; - const int32_t width = static_cast(tilingData_->width); - const int32_t historyCount = width - 1; - const int32_t ringStart = MAX_WIDTH - width; - const int32_t historyStartTok = tileStart - historyCount; - LocalTensor ring = inBuf.Get(); - bool hasGmHistoryCopy = false; - bool hasVectorInit = false; - const int64_t stateBaseOffset = static_cast(cacheIdx) * stateLen * dim + channelStart; - int64_t xHistoryOffset = static_cast(historyStartTok) * dim + channelStart; - - for (int32_t i = 0; i < ringStart; ++i) { - Duplicate(ring[i * MAX_BLOCK_DIM], static_cast(0), baseDim); - hasVectorInit = true; - } - - for (int32_t i = 0, srcTok = historyStartTok; i < historyCount; ++i, ++srcTok, xHistoryOffset += dim) { - LocalTensor histSlot = ring[(ringStart + i) * MAX_BLOCK_DIM]; - if (srcTok >= seqStart) { - DataCopy(histSlot, xGm[xHistoryOffset], baseDim); - hasGmHistoryCopy = true; - } else if (hasInit) { - const int32_t statePos = srcTok - seqStart + historyCount; - const int64_t stateOffset = stateBaseOffset + static_cast(statePos) * dim; - if (tilingData_->hasInitStateWorkspace != 0) { - const int64_t snapshotOffset = - (static_cast(seq) * historyCount + statePos) * dim + channelStart; - DataCopy(histSlot, initStateWorkspaceGm_[snapshotOffset], baseDim); - } else { - DataCopy(histSlot, convStatesGm[stateOffset], baseDim); - } - hasGmHistoryCopy = true; - } else { - Duplicate(histSlot, static_cast(0), baseDim); - hasVectorInit = true; - } - } - - if (hasGmHistoryCopy) { - SetFlag(stateMte2ToVEvent_); - WaitFlag(stateMte2ToVEvent_); - } - if (hasVectorInit) { - PipeBarrier(); - } - - if (tileLen > 0) { - const int32_t slot0 = SlotCurr(0); - const int64_t xOffset = static_cast(tileStart) * dim + channelStart; - DataCopy(ring[slot0 * MAX_BLOCK_DIM], xGm[xOffset], baseDim); - SetFlag(inputMte2ToVEvent_[slot0]); - } - - if (tileLen > 1) { - SetFlag(inputVToMte2Event_); - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessFnChunk(int32_t seq, int32_t cacheIdx, bool hasInit, - int32_t seqStart, int32_t seqLen, int32_t chunkStart, - int32_t chunkLen, int32_t channelStart, int32_t baseDim, - int32_t dim) -{ - LoadWeightAndBias(channelStart, baseDim); - InitRingSeqSplit(seq, cacheIdx, hasInit, seqStart, chunkStart, chunkLen, channelStart, baseDim, dim); - - RunSeq(chunkStart, chunkLen, channelStart, baseDim, dim); - - MaybeWriteBackSeqSplitTailChunk(chunkStart, chunkLen, seqStart, seqLen, cacheIdx, channelStart, baseDim, dim); - DrainTaskMte3(); -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::MaybeWriteBackSeqSplitTailChunk(int32_t chunkStart, int32_t chunkLen, - int32_t seqStart, int32_t seqLen, - int32_t cacheIdx, int32_t channelStart, - int32_t baseDim, int32_t dim) -{ - if (chunkStart + chunkLen != seqStart + seqLen) { - return; - } - - DrainTaskMte3(); - WriteBackState(cacheIdx, chunkLen, channelStart, baseDim, dim); -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::PrefetchInitStatesToWorkspace(int32_t channelStart, int32_t baseDimSize) -{ - if (tilingData_->hasInitStateWorkspace == 0) { - return; - } - - const int32_t dim = tilingData_->dim; - const int32_t historyCount = static_cast(tilingData_->width - 1); - const int32_t batch = tilingData_->batch; - const bool hasCacheIndices = (tilingData_->hasCacheIndices != 0); - const bool hasInitialStateMode = (tilingData_->hasInitialStateMode != 0); - LocalTensor tmpBuf = inBuf.Get()[0 * MAX_BLOCK_DIM]; - - for (int32_t seq = 0; seq < batch; ++seq) { - if (!ResolveSeqHasInit(seq, hasInitialStateMode)) { - continue; - } - - int32_t cacheIdx = 0; - if (!ResolveSeqCacheIndex(seq, hasCacheIndices, cacheIdx)) { - continue; - } - - const int64_t stateBaseOffset = static_cast(cacheIdx) * tilingData_->stateLen * dim + channelStart; - const int64_t snapshotBaseOffset = static_cast(seq) * historyCount * dim + channelStart; - for (int32_t statePos = 0; statePos < historyCount; ++statePos) { - const int64_t stateOffset = stateBaseOffset + static_cast(statePos) * dim; - const int64_t snapshotOffset = snapshotBaseOffset + static_cast(statePos) * dim; - DataCopy(tmpBuf, convStatesGm[stateOffset], baseDimSize); - SetFlag(initSnapshotMte2ToMte3Event_); - WaitFlag(initSnapshotMte2ToMte3Event_); - DataCopy(initStateWorkspaceGm_[snapshotOffset], tmpBuf, baseDimSize); - SetFlag(initSnapshotMte3ToMte2Event_); - WaitFlag(initSnapshotMte3ToMte2Event_); - } - } -} - -template -__aicore__ inline void CAUSAL_CONV1D_CLASS::ProcessVarlenTokenTiled() -{ - const int32_t dim = tilingData_->dim; - const int32_t batch = tilingData_->batch; - const int32_t seqLen = tilingData_->seqLen; - const int32_t cuSeqlen = tilingData_->cuSeqlen; - const int32_t baseDim = static_cast(tilingData_->baseDim); - const int32_t baseDimCnt = static_cast(tilingData_->baseDimCnt); - const int32_t tokenBlockSize = static_cast(tilingData_->tokenBlockSize); - const int32_t tokenBlockCnt = static_cast(tilingData_->tokenBlockCnt); - const bool hasCacheIndices = (tilingData_->hasCacheIndices != 0); - const bool hasInitialStateMode = (tilingData_->hasInitialStateMode != 0); - const bool isVarlenMode = (tilingData_->inputMode == 0); - - const int32_t blockIdx = static_cast(GetBlockIdx()); - const auto blockTask = - ResolveFnDirectBlockTask(blockIdx, tokenBlockCnt, tokenBlockSize, cuSeqlen, baseDimCnt, baseDim, dim); - if (tilingData_->hasInitStateWorkspace != 0) { - if (IsFnInitStateSnapshotOwnerBlock(blockTask)) { - PrefetchInitStatesToWorkspace(blockTask.channelStart, blockTask.baseDimSize); - } - SyncAll(); - } - if (!blockTask.valid) { - return; - } - - int32_t seq = 0; - int32_t seqUpperBound = batch; - if (isVarlenMode) { - if (!ResolveExplicitTokenTileSeqRange(blockTask.tokenTileId, seq, seqUpperBound)) { - seq = FindVarlenSeqByToken(blockTask.tokenStart); - } - } else { - seq = (seqLen > 0) ? (blockTask.tokenStart / seqLen) : 0; - } - - int32_t cursor = blockTask.tokenStart; - while (cursor < blockTask.tokenEnd && seq < seqUpperBound) { - int32_t seqStart = 0; - int32_t curSeqLen = 0; - if (!ResolveSeqTaskWindow(seq, tilingData_->inputMode, seqLen, seqStart, curSeqLen)) { - ++seq; - continue; - } - const int32_t curSeqEnd = seqStart + curSeqLen; - if (cursor < seqStart) { - cursor = seqStart; - } - if (cursor >= curSeqEnd) { - ++seq; - continue; - } - - const int32_t tileEnd = (blockTask.tokenEnd <= curSeqEnd) ? blockTask.tokenEnd : curSeqEnd; - const int32_t tileLen = tileEnd - cursor; - if (tileLen <= 0) { - ++seq; - continue; - } - - int32_t cacheIdx = 0; - if (!ResolveSeqCacheIndex(seq, hasCacheIndices, cacheIdx)) { - cursor = tileEnd; - ++seq; - continue; - } - - const bool hasInit = ResolveSeqHasInit(seq, hasInitialStateMode); - ProcessFnChunk(seq, cacheIdx, hasInit, seqStart, curSeqLen, cursor, tileLen, blockTask.channelStart, - blockTask.baseDimSize, dim); - - cursor = tileEnd; - ++seq; - } -} - -#endif // CUSTOM_CAUSAL_CONV1D_FN_TASKS_H diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h deleted file mode 100644 index 6360deb60..000000000 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_data.h +++ /dev/null @@ -1,66 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -/*! - * \file causal_conv1d_tiling_data.h - */ - -#ifndef CUSTOM_CAUSAL_CONV1D_TILING_DATA_H_ -#define CUSTOM_CAUSAL_CONV1D_TILING_DATA_H_ - -#include - -enum FnExecutionPlan : int64_t { - FN_EXECUTION_PLAN_INVALID = 0, - FN_EXECUTION_PLAN_CUTBS = 1, - FN_EXECUTION_PLAN_CUTBSD = 2, -}; - -inline constexpr int64_t ResolveFnExecutionPlan(int64_t baseDimCnt) -{ - return (baseDimCnt <= 0) ? FN_EXECUTION_PLAN_INVALID - : (baseDimCnt <= 1) ? FN_EXECUTION_PLAN_CUTBS - : FN_EXECUTION_PLAN_CUTBSD; -} - -struct CausalConv1dTilingData { - int64_t dim; - int64_t cuSeqlen; - int64_t seqLen; - int64_t inputMode; - - int64_t width; - - int64_t stateLen; - int64_t numCacheLines; - int64_t batch; - int64_t activationMode; - int64_t padSlotId; - int64_t hasBias; - int64_t baseDim; - int64_t baseDimCnt; - int64_t hasNumAcceptedTokens; - int64_t hasCacheIndices; - int64_t hasInitialStateMode; - int64_t tokenBlockSize; - int64_t tokenBlockCnt; - int64_t hasExplicitTokenSeqRanges; - int64_t explicitTokenSeqRangeCount; - int64_t tokenTileStartSeq[128]; - int64_t tokenTileEndSeq[128]; - int64_t hasInitStateWorkspace; - - int64_t dtypeKey; - int64_t runModeKey; - int64_t widthKey; - int64_t fnPlanKey; -}; -#endif // CUSTOM_CAUSAL_CONV1D_TILING_DATA_H_ diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_key.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_key.h deleted file mode 100644 index 773da5821..000000000 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d_tiling_key.h +++ /dev/null @@ -1,30 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -/*! - * \file causal_conv1d_tiling_key.h - * \brief causal_conv1d tiling key constants - */ - -#ifndef __CUSTOM_CAUSAL_CONV1D_TILING_KEY_H__ -#define __CUSTOM_CAUSAL_CONV1D_TILING_KEY_H__ - -#define CAUSAL_CONV1D_TPL_RUN_MODE_FN 0 -#define CAUSAL_CONV1D_TPL_RUN_MODE_UPDATE 1 -#define CAUSAL_CONV1D_TPL_WIDTH_RUNTIME 0 -#define CAUSAL_CONV1D_TPL_WIDTH_2 1 -#define CAUSAL_CONV1D_TPL_WIDTH_3 2 -#define CAUSAL_CONV1D_TPL_WIDTH_4 3 -#define CAUSAL_CONV1D_TPL_FN_PLAN_INVALID 0 -#define CAUSAL_CONV1D_TPL_FN_PLAN_CUTBS 1 -#define CAUSAL_CONV1D_TPL_FN_PLAN_CUTBSD 2 - -#endif // __CUSTOM_CAUSAL_CONV1D_TILING_KEY_H__ diff --git a/csrc/causal_conv1d/op_kernel/causal_conv1d_update.h b/csrc/causal_conv1d/op_kernel/causal_conv1d_update.h deleted file mode 100644 index 282de1adc..000000000 --- a/csrc/causal_conv1d/op_kernel/causal_conv1d_update.h +++ /dev/null @@ -1,73 +0,0 @@ -/** - * This program is free software, you can redistribute it and/or modify it. - * Copyright (c) 2025 Huawei Technologies Co., Ltd. - * This file is a part of the CANN Open Software. - * Licensed under CANN Open Software License Agreement Version 2.0 (the "License"). - * Please refer to the License for details. You may not use this file except in compliance with the License. - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING - * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. - * See LICENSE in the root of the software repository for the full text of the License. - */ - -#ifndef CUSTOM_CAUSAL_CONV1D_UPDATE_H -#define CUSTOM_CAUSAL_CONV1D_UPDATE_H - -#include "causal_conv1d.h" - -namespace NsCausalConv1d { - -template -class CausalConv1dUpdate : public CausalConv1d -{ -public: - __aicore__ inline void Init(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, GM_ADDR queryStartLoc, - GM_ADDR cacheIndices, GM_ADDR, GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, - const __gm__ CausalConv1dTilingData *tilingData) - { - (void)workspace; - this->ResetRuntimeState(tilingData); - this->xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(x)); - this->weightGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(weight)); - this->biasGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(bias)); - this->convStatesGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(convStates)); - this->queryStartLocGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(queryStartLoc)); - this->cacheIndicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(cacheIndices)); - this->numAcceptedTokensGm.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(numAcceptedTokens)); - this->yGm.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(y)); - this->InitSharedBuffersAndEvents(); - } - - __aicore__ inline void Process() - { - const __gm__ CausalConv1dTilingData *tilingData = this->GetTilingData(); - const int32_t dim = tilingData->dim; - const int32_t baseDimCnt = static_cast(tilingData->baseDimCnt); - const int32_t width = static_cast(tilingData->width); - const int32_t baseDim = static_cast(tilingData->baseDim); - if (baseDim <= 0 || baseDimCnt <= 0 || baseDim > MAX_BLOCK_DIM || width < 2 || width > MAX_WIDTH || dim <= 0 || - tilingData->batch <= 0) { - this->ReleaseEvents(); - return; - } - - this->ProcessDefault(); - this->ReleaseEvents(); - } -}; - -template -__aicore__ inline void RunCausalConv1dUpdate(GM_ADDR x, GM_ADDR weight, GM_ADDR bias, GM_ADDR convStates, - GM_ADDR queryStartLoc, GM_ADDR cacheIndices, GM_ADDR initialStateMode, - GM_ADDR numAcceptedTokens, GM_ADDR y, GM_ADDR workspace, - const __gm__ CausalConv1dTilingData *tilingData) -{ - CausalConv1dUpdate op; - op.Init(x, weight, bias, convStates, queryStartLoc, cacheIndices, initialStateMode, numAcceptedTokens, y, workspace, - tilingData); - op.Process(); -} - -} // namespace NsCausalConv1d - -#endif // CUSTOM_CAUSAL_CONV1D_UPDATE_H diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp index 1cfccb0d7..bc5e5c432 100644 --- a/csrc/pytorch_extensions.cpp +++ b/csrc/pytorch_extensions.cpp @@ -143,10 +143,9 @@ TORCH_LIBRARY_FRAGMENT(npu, m) "Tensor? query_start_loc=None, bool activation_mode=False, int pad_slot_id=-1) -> Tensor"); m.def( - "causal_conv1d(Tensor x, Tensor weight, Tensor conv_states, Tensor? bias=None, " - "Tensor? query_start_loc=None, Tensor? cache_indices=None, Tensor? has_initial_state=None, " - "Tensor? num_accepted_tokens=None, int activation_mode=0, int pad_slot_id=-1, " - "int run_mode=0) -> Tensor"); + "causal_conv1d(Tensor x, Tensor weight, Tensor conv_states, " + "Tensor query_start_loc, Tensor cache_indices, Tensor has_initial_state, " + "Tensor? bias=None, bool activation_mode=False, int pad_slot_id=-1) -> Tensor"); } } // namespace @@ -226,25 +225,15 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) }); m.impl("causal_conv1d", [](const at::Tensor &x, const at::Tensor &weight, const at::Tensor &conv_states, - const c10::optional &bias, const c10::optional &query_start_loc, - const c10::optional &cache_indices, - const c10::optional &has_initial_state, - const c10::optional &num_accepted_tokens, int64_t activation_mode, - int64_t pad_slot_id, int64_t run_mode) { - // Handle optional parameters - convert None to empty tensors - auto bias_or_empty = bias.has_value() ? *bias : at::empty({0}, x.options()); - auto query_start_loc_or_empty = - query_start_loc.has_value() ? *query_start_loc : at::empty({0}, x.options().dtype(at::kLong)); - auto cache_indices_or_empty = - cache_indices.has_value() ? *cache_indices : at::empty({0}, x.options().dtype(at::kLong)); - auto has_initial_state_or_empty = - has_initial_state.has_value() ? *has_initial_state : at::empty({0}, x.options().dtype(at::kLong)); - auto num_accepted_tokens_or_empty = - num_accepted_tokens.has_value() ? *num_accepted_tokens : at::empty({0}, x.options().dtype(at::kLong)); - - return sglang::npu_kernel::causal_conv1d_impl( - x, weight, bias_or_empty, conv_states, query_start_loc_or_empty, cache_indices_or_empty, - has_initial_state_or_empty, num_accepted_tokens_or_empty, activation_mode, pad_slot_id, run_mode); + const at::Tensor &query_start_loc, const at::Tensor &cache_indices, + const at::Tensor &has_initial_state, const c10::optional &bias, + bool activation_mode, int64_t pad_slot_id) { + // Cast PyTorch-default int64/bool to the kernel's int32/bool and make bias + // contiguous (all no-ops when already correct) to satisfy the strict TORCH_CHECKs. + auto bias_or_empty = bias.has_value() ? bias->contiguous() : at::empty({0}, x.options()); + return sglang::npu_kernel::causal_conv1d_impl(x, weight, conv_states, query_start_loc.to(at::kInt), + cache_indices.to(at::kInt), has_initial_state.to(at::kBool), + bias_or_empty, activation_mode, pad_slot_id); }); } } // namespace diff --git a/tests/python/sgl_kernel_npu/test_conv1d_prefill.py b/tests/python/sgl_kernel_npu/test_conv1d_prefill.py index 9b18ab228..59366cead 100644 --- a/tests/python/sgl_kernel_npu/test_conv1d_prefill.py +++ b/tests/python/sgl_kernel_npu/test_conv1d_prefill.py @@ -365,18 +365,65 @@ def run_negative_cases(device: torch.device, dtype: torch.dtype, pad_slot_id: in ("dtype must match",), ) + # [7] wrapper casts int64/bool -> int32/bool, so PyTorch-default dtypes are accepted + # (int64 query_start_loc used to be rejected). + y_default_dtypes = torch.ops.npu.causal_conv1d( + x, + weight, + conv_states, + make_device_long_tensor([0, 4, 8], device), # int64 query_start_loc + make_device_long_tensor([0, 3], device), # int64 cache_indices + make_device_long_tensor([1, 0], device), # int64 "has_initial_state" -> bool + bias=bias, + ) + torch.npu.synchronize() + assert ( + y_default_dtypes.shape == x.shape + ), "int64/bool inputs should be accepted (wrapper casts)" + print("[PASS] accepts_default_int64_bool_inputs (wrapper casts to int32/bool)") + + # [5] varlen: an empty/too-short query_start_loc would underflow batch -> must reject. + x2d = torch.randn((4, dim), device=device, dtype=dtype) expect_failure( - "dtype_mismatch_query_start_loc", + "empty_query_start_loc_varlen", lambda: torch.ops.npu.causal_conv1d( - x, + x2d, weight, conv_states, - make_device_long_tensor([0, 4, 8], device), + make_device_int_tensor([0], device), # size 1 (< 2) cache_indices, has_initial_state, bias=bias, ), - ("query_start_loc dtype must be int32",), + ("query_start_loc", "at least 2 elements"), + ) + + # [6] cache_indices / has_initial_state are indexed per sequence -> size must be >= batch (=2). + expect_failure( + "cache_indices_too_small", + lambda: torch.ops.npu.causal_conv1d( + x, + weight, + conv_states, + query_start_loc, + make_device_int_tensor([0], device), # size 1 (< batch 2) + has_initial_state, + bias=bias, + ), + ("cache_indices", "size >= batch"), + ) + expect_failure( + "has_initial_state_too_small", + lambda: torch.ops.npu.causal_conv1d( + x, + weight, + conv_states, + query_start_loc, + cache_indices, + make_device_bool_tensor([True], device), # size 1 (< batch 2) + bias=bias, + ), + ("has_initial_state", "size >= batch"), )