From 2b0da3cbb59071905ea2411d7c12610d37458e9a Mon Sep 17 00:00:00 2001 From: Ringssss <28751159+Ringssss@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:59:09 +0800 Subject: [PATCH] feat(hygon): add native Clover diffusion and sparse MoE operators --- CMakeLists.txt | 6 +- README.md | 3 + THIRD_PARTY.txt | 8 + docs/clover-hygon.md | 77 +++++++ docs/operator-api-alignment.md | 6 + scripts/generate_wrappers.py | 11 + src/CMakeLists.txt | 8 + src/base/clover_block_conv1d.h | 49 +++++ src/base/clover_coreset.h | 23 ++ src/base/clover_live_mask.h | 22 ++ src/base/clover_local_memory.h | 59 +++++ src/base/clover_mask_routes.h | 22 ++ src/base/clover_merge_moe.h | 24 +++ src/base/clover_recurrent_gdn.h | 51 +++++ src/native/cuda/hygon/caster.cuh | 3 + src/native/cuda/hygon/ops/cast/kernel.h | 192 +++++++++++++++++ .../cuda/hygon/ops/causal_conv1d/kernel.h | 202 ++++++++++++++++++ src/native/cuda/hygon/ops/fill/kernel.h | 16 ++ .../ops/fused_gated_delta_net_gating/kernel.h | 120 +++++++++++ .../hygon/ops/moe_align_block_size/kernel.cu | 81 +++++++ .../hygon/ops/moe_align_block_size/kernel.cuh | 86 ++++++++ .../hygon/ops/moe_align_block_size/kernel.h | 35 +++ src/native/cuda/hygon/ops/mrope/kernel.h | 12 ++ src/native/cuda/hygon/ops/mul/kernel.h | 16 ++ .../hygon/ops/separate_table_rope/kernel.h | 27 +++ src/native/cuda/hygon/ops/tril/kernel.h | 20 ++ src/native/cuda/hygon/ops/triu/kernel.h | 20 ++ src/triton/jit/backend.h | 16 +- src/triton/jit/compile.py | 35 ++- src/triton/jit/compiler.cc | 17 ++ src/triton/ops/add/jit.h | 7 + src/triton/ops/clover_block_conv1d/jit.h | 56 +++++ src/triton/ops/clover_block_conv1d/jit.py | 145 +++++++++++++ src/triton/ops/clover_coreset/jit.h | 46 ++++ src/triton/ops/clover_coreset/jit.py | 52 +++++ src/triton/ops/clover_live_mask/jit.h | 34 +++ src/triton/ops/clover_live_rows/jit.py | 39 ++++ src/triton/ops/clover_local_memory/jit.h | 62 ++++++ src/triton/ops/clover_local_memory/jit.py | 188 ++++++++++++++++ src/triton/ops/clover_mask_routes/jit.h | 34 +++ src/triton/ops/clover_merge_moe/jit.h | 35 +++ src/triton/ops/clover_recurrent_gdn/jit.h | 51 +++++ src/triton/ops/clover_recurrent_gdn/jit.py | 77 +++++++ src/triton/ops/moe_fused_dense/jit.h | 93 ++++++++ src/triton/ops/moe_fused_dense/jit.py | 103 +++++++++ tests/conftest.py | 2 +- tests/native/build_clover_workspace.py | 37 ++++ tests/native/check_clover_workspace.py | 68 ++++++ tests/native/clover_workspace_smoke.cpp | 41 ++++ tests/test_clover_live_rows.py | 30 +++ tests/test_clover_preview.py | 126 +++++++++++ tests/test_generate_wrappers.py | 12 ++ tests/test_triton_jit_compile.py | 24 +++ 53 files changed, 2620 insertions(+), 9 deletions(-) create mode 100644 docs/clover-hygon.md create mode 100644 src/base/clover_block_conv1d.h create mode 100644 src/base/clover_coreset.h create mode 100644 src/base/clover_live_mask.h create mode 100644 src/base/clover_local_memory.h create mode 100644 src/base/clover_mask_routes.h create mode 100644 src/base/clover_merge_moe.h create mode 100644 src/base/clover_recurrent_gdn.h create mode 100644 src/native/cuda/hygon/ops/cast/kernel.h create mode 100644 src/native/cuda/hygon/ops/causal_conv1d/kernel.h create mode 100644 src/native/cuda/hygon/ops/fill/kernel.h create mode 100644 src/native/cuda/hygon/ops/fused_gated_delta_net_gating/kernel.h create mode 100644 src/native/cuda/hygon/ops/moe_align_block_size/kernel.cu create mode 100644 src/native/cuda/hygon/ops/moe_align_block_size/kernel.cuh create mode 100644 src/native/cuda/hygon/ops/moe_align_block_size/kernel.h create mode 100644 src/native/cuda/hygon/ops/mrope/kernel.h create mode 100644 src/native/cuda/hygon/ops/mul/kernel.h create mode 100644 src/native/cuda/hygon/ops/separate_table_rope/kernel.h create mode 100644 src/native/cuda/hygon/ops/tril/kernel.h create mode 100644 src/native/cuda/hygon/ops/triu/kernel.h create mode 100644 src/triton/ops/clover_block_conv1d/jit.h create mode 100644 src/triton/ops/clover_block_conv1d/jit.py create mode 100644 src/triton/ops/clover_coreset/jit.h create mode 100644 src/triton/ops/clover_coreset/jit.py create mode 100644 src/triton/ops/clover_live_mask/jit.h create mode 100644 src/triton/ops/clover_live_rows/jit.py create mode 100644 src/triton/ops/clover_local_memory/jit.h create mode 100644 src/triton/ops/clover_local_memory/jit.py create mode 100644 src/triton/ops/clover_mask_routes/jit.h create mode 100644 src/triton/ops/clover_merge_moe/jit.h create mode 100644 src/triton/ops/clover_recurrent_gdn/jit.h create mode 100644 src/triton/ops/clover_recurrent_gdn/jit.py create mode 100644 src/triton/ops/moe_fused_dense/jit.h create mode 100644 src/triton/ops/moe_fused_dense/jit.py create mode 100644 tests/native/build_clover_workspace.py create mode 100644 tests/native/check_clover_workspace.py create mode 100644 tests/native/clover_workspace_smoke.cpp create mode 100644 tests/test_clover_live_rows.py create mode 100644 tests/test_clover_preview.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 02c416958..ac1bb0c78 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,6 +31,8 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) set(PYBIND11_ENABLE_EXTRAS ON) # Options for backends. +set(INFINI_OPS_DECLARATIONS "" CACHE STRING "Additional public operator APIs, even without active providers") + option(WITH_CPU "Enable CPU backend" OFF) option(WITH_NVIDIA "Enable CUDA backend" OFF) option(WITH_ILUVATAR "Enable Iluvatar GPU backend" OFF) @@ -419,10 +421,10 @@ if(WITH_NINETOOTHED) endif() if(WITH_TRITON) - if(NOT WITH_NVIDIA) + if(NOT WITH_NVIDIA AND NOT WITH_HYGON) message( FATAL_ERROR - "`WITH_TRITON` requires `WITH_NVIDIA=ON` because NVIDIA is the only implemented Triton JIT backend." + "`WITH_TRITON` requires `WITH_NVIDIA=ON` or `WITH_HYGON=ON`." ) endif() if(NOT GENERATE_PYTHON_BINDINGS) diff --git a/README.md b/README.md index 0d784b317..c0bf722f3 100644 --- a/README.md +++ b/README.md @@ -80,3 +80,6 @@ Useful entry points: ## License See [LICENSE](LICENSE). + +Hygon Clover diffusion and native Triton JIT setup are described in +[the Clover operator guide](docs/clover-hygon.md). diff --git a/THIRD_PARTY.txt b/THIRD_PARTY.txt index e588b1443..3cc8a0337 100644 --- a/THIRD_PARTY.txt +++ b/THIRD_PARTY.txt @@ -113,3 +113,11 @@ from actual package metadata. That conservative package inventory can include materials for components outside these two DSOs; it is intentionally not filtered by fixed hashes or used to restrict build capability. A complete wheel or SDK redistribution must retain its own applicable notices, including secondary ones. + +Clover preview kernels +src/triton/ops/clover_block_conv1d/jit.py and +src/triton/ops/clover_local_memory/jit.py are adapted from SGLang PR #46927, +WilliamZR/sglang commit 058939af24d5cdc44bb4b8852d0ddfab555ed71e, +python/sglang/kernels/ops/{mamba/clover_conv_triton.py,attention/clover_gdn_triton.py}. +They retain Apache-2.0 licensing. Changes provide InfiniOps JIT entry points, +Hygon launch geometry, and bounds checks for the state pool. diff --git a/docs/clover-hygon.md b/docs/clover-hygon.md new file mode 100644 index 000000000..b17a51db1 --- /dev/null +++ b/docs/clover-hygon.md @@ -0,0 +1,77 @@ +# Hygon Clover operators + +These operators support the companion InfiniLM Clover implementation. Model +forward stays in C++; Hygon driver modules launch warmed Triton binaries. +Compilation uses the installed DTK Triton HIP target (for example `gfx936`, +wave size 64), not CUDA compute-capability arithmetic. + +## Build + +Install an InfiniRT revision containing the Hygon driver interface first +([InfiniRT PR #52](https://github.com/InfiniTensor/InfiniRT/pull/52)). Then: + +```sh +cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_INSTALL_PREFIX="$INFINI_ROOT" -DINFINI_RT_ROOT="$INFINI_ROOT" \ + -DWITH_CPU=ON -DWITH_HYGON=ON -DHYGON_ARCH=gfx936 -DWITH_TRITON=ON \ + -DGENERATE_PYTHON_BINDINGS=ON -DAUTO_DETECT_DEVICES=OFF \ + -DAUTO_DETECT_BACKENDS=OFF +cmake --build build -j 12 +cmake --install build +``` + +The Python `infini.triton` package must be importable. CMake installs its package +contents at the prefix, while wheel installation supplies the `infini` namespace. +Include DTK's CUDA compatibility library directory in `LD_LIBRARY_PATH`. + +`INFINI_OPS_DECLARATIONS` accepts a semicolon-separated list of additional public +operator APIs. These can have no provider in the selected build: capability +queries return unsupported, and execution is rejected. This allows a downstream +application to compile its complete model registry without claiming that every +registered model runs on Hygon. InfiniLM supplies a matching declaration preset. + +## Operator contracts and source alignment + +The Clover-specific APIs target the released Clover model equations and state +layout, rather than claiming to be PyTorch standard operators. Preview kernel +provenance and licensing are recorded in `THIRD_PARTY.txt`. + +| Operator | Inputs and behavior | +| --- | --- | +| `CloverBlockConv1d` | Packed `[tokens,channels]` BF16/FP16 input, `[slots,channels,3]` clean history, slot indices, left `[channels,1,4]` and learned right `[channels,1,3]` weights, int8 validity, block width up to 32. Produces SiLU output without updating history. | +| `CloverLocalMemory` | Q/K/V `[1,tokens,heads,128]`, FP32 gates and `[slots,value_heads,128,128]` state, slot indices, validity and three per-head scalars. Implements bidirectional block attention plus the decayed clean-history bucket. Does not update state. | +| `CloverRecurrentGdn` | Packed Q/K/V, gates/beta, sequence offsets and initial/final slot vectors. Implements normalized causal delta recurrence with FP32 persistent state; writes sequence output and final state. Input/final slot indices must be valid and final destinations must not race. | +| `CloverCoreset` | Contiguous `[tokens,256]` router logits, block width and routing mass. Keeps the stable top-p support of aggregated per-token top-8 scores and masks other logits. Requires an explicit workspace. | +| `CloverLiveMask` | Current/previous int64 token IDs and the mask ID. Marks rows live if either token is masked, or refresh is forced; updates previous IDs. | +| `CloverMaskRoutes` | Replaces inactive rows' int32 expert IDs with the out-of-range sentinel `num_experts`. | +| `CloverMergeMoe` | Selects fresh routed output for live rows and cached output for other rows; updates the cache. Output and fresh input may alias. | +| `MoeFusedDense` (Hygon) | Existing packed dense-MoE API, with BF16/FP16 weights, top-k weights/IDs and aligned expert metadata. Triton computes gate/up, fused SiLU/down, and weighted reduction. Invalid routes contribute zero; alignment block size 16 is supported and expert biases are unsupported. Requires an explicit workspace. | + +The live-row protocol is an optional approximation for iterative diffusion, +adapted from the supplied dllm-h100 TV6 study. It is not a general guarantee that +fixed token IDs imply unchanged hidden states. InfiniLM refreshes every five +previews and always runs full causal replay before committing state. + +The standard `MoeAlignBlockSize`, causal convolution, gating, cast, fill, +multiplication, triangular masks and rotary providers also gain Hygon coverage. +They reuse this repository's existing native CUDA-like implementations. + +## Validation + +Actual BW/DTK runs compare preview operators to independent Torch equations and +verify bitwise state immutability, mixed-length causal state commits, coreset +selection, dense MoE, liveness and cache merging. Public-header and wrapper/JIT +unit tests cover the HIP architecture string and metadata without CUDA scratch +fields. The companion InfiniLM campaign runs full checkpoint generation. + +Workspace-bearing operators are invoked through the C++ API with an owning +workspace. The native test adapter in `tests/native/clover_workspace_smoke.cpp` +exercises that path. A missing workspace is rejected before launching a kernel. + +Build and run the workspace adapter against the installed SDK: + +```sh +python tests/native/build_clover_workspace.py +PYTHONPATH="build/clover-tests:$PYTHONPATH" python tests/native/check_clover_workspace.py +python -m pytest tests/test_clover_preview.py tests/test_clover_live_rows.py -q +``` diff --git a/docs/operator-api-alignment.md b/docs/operator-api-alignment.md index 1118420ab..c70da2387 100644 --- a/docs/operator-api-alignment.md +++ b/docs/operator-api-alignment.md @@ -238,3 +238,9 @@ unless callers need it. When applying these rules, record the selected upstream interface and any intentional C++ adaptations in the pull request. Follow `CONTRIBUTING.md` for compatibility, testing, and validation requirements. + +## Clover diffusion state operators + +The `Clover*` operators use the released Clover model equations, with packed +request/state-pool arguments for inference. Their exact contracts and source +provenance are documented in [the Hygon Clover guide](clover-hygon.md). diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py index 38cb3e8f9..ab20b88bd 100644 --- a/scripts/generate_wrappers.py +++ b/scripts/generate_wrappers.py @@ -2293,6 +2293,12 @@ def _dispatch_gen_batch_size(): help="Include Triton backend implementations.", ) + parser.add_argument( + "--declarations", + nargs="*", + default=[], + help="Also emit operator APIs with no active provider.", + ) args = parser.parse_args() for directory in (_BINDINGS_DIR, _GENERATED_SRC_DIR, _INCLUDE_DIR): @@ -2309,6 +2315,11 @@ def _dispatch_gen_batch_size(): with_native_w4=args.with_native_w4, ) + for name in _normalize_op_allowlist(args.declarations): + if not (_BASE_DIR / f"{name}.h").is_file(): + raise ValueError(f"Unknown operator declaration: {name}") + ops.setdefault(name, []) + if config_path is not None: config = ops_config.load_ops_config(config_path) ops = _select_ops_from_config(ops, config, config_path) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6c6f3710c..eb04f3e2e 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -249,6 +249,11 @@ if(WITH_TRITON OR GENERATE_PYTHON_BINDINGS) find_package(pybind11 CONFIG REQUIRED) endif() +if(WITH_TRITON AND WITH_HYGON) + target_link_directories(infiniops PRIVATE "${HYGON_CUDA_ROOT}/lib64") + target_link_libraries(infiniops PRIVATE cuda) +endif() + if(WITH_TRITON) target_link_libraries(infiniops PRIVATE Python::Python @@ -1290,6 +1295,9 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS) # failures. set(GENERATOR_ARGS --devices ${DEVICE_LIST}) + if(INFINI_OPS_DECLARATIONS) + list(APPEND GENERATOR_ARGS --declarations ${INFINI_OPS_DECLARATIONS}) + endif() if(_infini_ops_ops_config) list(APPEND GENERATOR_ARGS --ops-config "${_infini_ops_ops_config}") diff --git a/src/base/clover_block_conv1d.h b/src/base/clover_block_conv1d.h new file mode 100644 index 000000000..c6b90fd3c --- /dev/null +++ b/src/base/clover_block_conv1d.h @@ -0,0 +1,49 @@ +#pragma once + +#include "operator.h" + +namespace infini::ops { + +// Clover preview convolution reads a clean prefix without updating its state. +class CloverBlockConv1d : public Operator { + public: + CloverBlockConv1d(const Tensor input, const Tensor state, + const Tensor indices, const Tensor weight, + const Tensor right_weight, const Tensor valid, + size_t block_size, Tensor out) { + assert((input.ndim() == 2 && out.shape() == input.shape() && + state.ndim() == 3 && weight.ndim() == 3 && + weight.shape() == Tensor::Shape{input.size(1), 1, 4} && + right_weight.shape() == Tensor::Shape{input.size(1), 1, 3} && + state.size(1) == input.size(1) && state.size(2) == 3 && + indices.ndim() == 1 && block_size > 0 && block_size <= 32 && + input.size(0) == indices.numel() * block_size && + valid.shape() == Tensor::Shape{input.size(0)}) && + "CloverBlockConv1d: invalid block or state shape."); + assert((input.dtype() == DataType::kBFloat16 || + input.dtype() == DataType::kFloat16) && + "CloverBlockConv1d: expected BF16 or FP16 input."); + for (const auto& tensor : {state, weight, right_weight, out}) { + assert(tensor.dtype() == input.dtype() && + "CloverBlockConv1d: input and weight dtypes must match."); + } + assert((indices.dtype() == DataType::kInt32 || + indices.dtype() == DataType::kInt64) && + indices.IsContiguous() && valid.dtype() == DataType::kInt8 && + valid.IsContiguous() && + "CloverBlockConv1d: expected contiguous indices and int8 validity."); + for (const auto& tensor : + {state, indices, weight, right_weight, valid, out}) { + assert(tensor.device() == input.device() && + "CloverBlockConv1d: device mismatch."); + } + assert(!out.HasBroadcastDim() && "CloverBlockConv1d: overlapping output."); + } + + virtual void operator()(const Tensor input, const Tensor state, + const Tensor indices, const Tensor weight, + const Tensor right_weight, const Tensor valid, + size_t block_size, Tensor out) const = 0; +}; + +} // namespace infini::ops diff --git a/src/base/clover_coreset.h b/src/base/clover_coreset.h new file mode 100644 index 000000000..198890fc9 --- /dev/null +++ b/src/base/clover_coreset.h @@ -0,0 +1,23 @@ +#pragma once + +#include "operator.h" + +namespace infini::ops { + +class CloverCoreset : public Operator { + public: + CloverCoreset(const Tensor logits, size_t block_size, double top_p, + Tensor out) { + assert(logits.ndim() == 2 && logits.size(1) == 256 && block_size > 0 && + block_size <= 32 && logits.size(0) % block_size == 0 && + out.shape() == logits.shape() && out.dtype() == logits.dtype() && + logits.IsContiguous() && out.IsContiguous() && + logits.device() == out.device() && top_p > 0 && top_p <= 1 && + "CloverCoreset: invalid logits, block size or routing mass."); + } + + virtual void operator()(const Tensor logits, size_t block_size, double top_p, + Tensor out) const = 0; +}; + +} // namespace infini::ops diff --git a/src/base/clover_live_mask.h b/src/base/clover_live_mask.h new file mode 100644 index 000000000..f688402db --- /dev/null +++ b/src/base/clover_live_mask.h @@ -0,0 +1,22 @@ +#pragma once +#include "operator.h" + +namespace infini::ops { + +class CloverLiveMask : public Operator { + public: + CloverLiveMask(const Tensor ids, Tensor previous, int64_t mask_id, + bool refresh, Tensor live) { + assert(ids.IsContiguous() && previous.IsContiguous() && + live.IsContiguous() && ids.numel() == previous.numel() && + ids.numel() == live.numel() && ids.dtype() == DataType::kInt64 && + previous.dtype() == ids.dtype() && live.dtype() == DataType::kInt8 && + ids.device() == previous.device() && ids.device() == live.device() && + "CloverLiveMask: invalid token state."); + } + + virtual void operator()(const Tensor ids, Tensor previous, int64_t mask_id, + bool refresh, Tensor live) const = 0; +}; + +} // namespace infini::ops diff --git a/src/base/clover_local_memory.h b/src/base/clover_local_memory.h new file mode 100644 index 000000000..f58757891 --- /dev/null +++ b/src/base/clover_local_memory.h @@ -0,0 +1,59 @@ +#pragma once + +#include "operator.h" + +namespace infini::ops { + +class CloverLocalMemory : public Operator { + public: + CloverLocalMemory(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor state, const Tensor indices, + const Tensor valid, const Tensor logit_scale, + const Tensor history_bias, const Tensor history_c, + size_t block_size, Tensor out) { + assert((q.ndim() == 4 && q.size(0) == 1 && k.shape() == q.shape() && + v.ndim() == 4 && v.size(0) == 1 && v.size(1) == q.size(1) && + q.size(2) > 0 && v.size(2) > 0 && v.size(2) % q.size(2) == 0 && + q.size(3) == 128 && v.size(3) == 128 && out.shape() == v.shape() && + state.ndim() == 4 && state.size(1) == v.size(2) && + state.size(2) == 128 && state.size(3) == 128 && + indices.ndim() == 1 && block_size > 0 && block_size <= 32 && + q.size(1) == indices.numel() * block_size && + g.shape() == Tensor::Shape{1, q.size(1), v.size(2)} && + valid.shape() == Tensor::Shape{q.size(1)}) && + "CloverLocalMemory: invalid sequence, head or state shape."); + assert( + (q.dtype() == DataType::kBFloat16 || q.dtype() == DataType::kFloat16) && + k.dtype() == q.dtype() && v.dtype() == q.dtype() && + out.dtype() == q.dtype() && state.dtype() == DataType::kFloat32 && + g.dtype() == DataType::kFloat32 && + "CloverLocalMemory: expected reduced data and FP32 state/gates."); + assert((indices.dtype() == DataType::kInt32 || + indices.dtype() == DataType::kInt64) && + indices.IsContiguous() && valid.dtype() == DataType::kInt8 && + valid.IsContiguous() && + "CloverLocalMemory: invalid index/validity dtype."); + for (const auto& scalar : {logit_scale, history_bias, history_c}) { + assert(scalar.ndim() == 1 && scalar.numel() == v.size(2) && + scalar.IsContiguous() && + (scalar.dtype() == q.dtype() || + scalar.dtype() == DataType::kFloat32) && + "CloverLocalMemory: invalid per-head scalars."); + } + for (const auto& tensor : {k, v, g, state, indices, valid, logit_scale, + history_bias, history_c, out}) { + assert(tensor.device() == q.device() && + "CloverLocalMemory: device mismatch."); + } + assert(!out.HasBroadcastDim() && "CloverLocalMemory: overlapping output."); + } + + virtual void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor state, + const Tensor indices, const Tensor valid, + const Tensor logit_scale, const Tensor history_bias, + const Tensor history_c, size_t block_size, + Tensor out) const = 0; +}; + +} // namespace infini::ops diff --git a/src/base/clover_mask_routes.h b/src/base/clover_mask_routes.h new file mode 100644 index 000000000..d04b5ad99 --- /dev/null +++ b/src/base/clover_mask_routes.h @@ -0,0 +1,22 @@ +#pragma once +#include "operator.h" + +namespace infini::ops { + +class CloverMaskRoutes : public Operator { + public: + CloverMaskRoutes(const Tensor ids, const Tensor live, int64_t num_experts, + Tensor out) { + assert(ids.ndim() == 2 && ids.IsContiguous() && live.IsContiguous() && + out.IsContiguous() && live.numel() == ids.size(0) && + out.shape() == ids.shape() && ids.dtype() == DataType::kInt32 && + out.dtype() == ids.dtype() && live.dtype() == DataType::kInt8 && + ids.device() == live.device() && ids.device() == out.device() && + "CloverMaskRoutes: invalid routes."); + } + + virtual void operator()(const Tensor ids, const Tensor live, + int64_t num_experts, Tensor out) const = 0; +}; + +} // namespace infini::ops diff --git a/src/base/clover_merge_moe.h b/src/base/clover_merge_moe.h new file mode 100644 index 000000000..6ccc62a81 --- /dev/null +++ b/src/base/clover_merge_moe.h @@ -0,0 +1,24 @@ +#pragma once +#include "operator.h" + +namespace infini::ops { + +class CloverMergeMoe : public Operator { + public: + CloverMergeMoe(const Tensor fresh, const Tensor live, Tensor cache, + Tensor out) { + assert(fresh.ndim() == 2 && live.numel() == fresh.size(0) && + cache.shape() == fresh.shape() && out.shape() == fresh.shape() && + cache.dtype() == fresh.dtype() && out.dtype() == fresh.dtype() && + live.dtype() == DataType::kInt8 && fresh.IsContiguous() && + cache.IsContiguous() && out.IsContiguous() && live.IsContiguous() && + fresh.device() == live.device() && + fresh.device() == cache.device() && fresh.device() == out.device() && + "CloverMergeMoe: invalid cached rows."); + } + + virtual void operator()(const Tensor fresh, const Tensor live, Tensor cache, + Tensor out) const = 0; +}; + +} // namespace infini::ops diff --git a/src/base/clover_recurrent_gdn.h b/src/base/clover_recurrent_gdn.h new file mode 100644 index 000000000..10dd76066 --- /dev/null +++ b/src/base/clover_recurrent_gdn.h @@ -0,0 +1,51 @@ +#pragma once + +#include "operator.h" + +namespace infini::ops { + +class CloverRecurrentGdn : public Operator { + public: + CloverRecurrentGdn(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor beta, const Tensor offsets, + const Tensor initial_indices, const Tensor final_indices, + Tensor state, Tensor out) { + assert((q.ndim() == 4 && q.size(0) == 1 && k.shape() == q.shape() && + v.ndim() == 4 && v.size(0) == 1 && q.size(1) == v.size(1) && + q.size(2) > 0 && v.size(2) % q.size(2) == 0 && q.size(3) == 128 && + v.size(3) == 128 && out.shape() == v.shape() && + g.shape() == Tensor::Shape{1, q.size(1), v.size(2)} && + beta.shape() == g.shape() && state.ndim() == 4 && + state.size(1) == v.size(2) && state.size(2) == 128 && + state.size(3) == 128 && offsets.ndim() == 1 && + offsets.numel() >= 2 && initial_indices.ndim() == 1 && + final_indices.shape() == initial_indices.shape() && + initial_indices.numel() + 1 == offsets.numel()) && + "CloverRecurrentGdn: invalid sequence or state shape."); + assert( + (q.dtype() == DataType::kBFloat16 || q.dtype() == DataType::kFloat16) && + k.dtype() == q.dtype() && v.dtype() == q.dtype() && + out.dtype() == q.dtype() && state.dtype() == DataType::kFloat32 && + g.dtype() == DataType::kFloat32 && beta.dtype() == DataType::kFloat32 && + "CloverRecurrentGdn: invalid dtype."); + for (const auto& index : {offsets, initial_indices, final_indices}) { + assert(index.IsContiguous() && + (index.dtype() == DataType::kInt32 || + index.dtype() == DataType::kInt64) && + "CloverRecurrentGdn: invalid sequence indices."); + } + for (const auto& tensor : + {k, v, g, beta, offsets, initial_indices, final_indices, state, out}) { + assert(tensor.device() == q.device() && + "CloverRecurrentGdn: device mismatch."); + } + } + + virtual void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor beta, + const Tensor offsets, const Tensor initial_indices, + const Tensor final_indices, Tensor state, + Tensor out) const = 0; +}; + +} // namespace infini::ops diff --git a/src/native/cuda/hygon/caster.cuh b/src/native/cuda/hygon/caster.cuh index 4b381d03b..901fc1fd5 100644 --- a/src/native/cuda/hygon/caster.cuh +++ b/src/native/cuda/hygon/caster.cuh @@ -7,6 +7,9 @@ namespace infini::ops { +using infini::rt::cuda_bfloat16; +using infini::rt::cuda_bfloat162; + namespace detail { template <> diff --git a/src/native/cuda/hygon/ops/cast/kernel.h b/src/native/cuda/hygon/ops/cast/kernel.h new file mode 100644 index 000000000..250195f06 --- /dev/null +++ b/src/native/cuda/hygon/ops/cast/kernel.h @@ -0,0 +1,192 @@ +#ifndef INFINI_OPS_HYGON_CAST_KERNEL_H_ +#define INFINI_OPS_HYGON_CAST_KERNEL_H_ + +#include + +#include +#include +#include +#include +#include +#include + +#include "base/cast.h" +#include "native/cuda/hygon/caster.cuh" + +namespace infini::ops { +namespace hygon_cast_detail { +struct Layout { + size_t rank{}; + size_t shape[16]{}; + ptrdiff_t input_strides[16]{}; + ptrdiff_t output_strides[16]{}; + bool contiguous{}; +}; +template +__global__ void Convert(const InT* input, OutT* out, std::size_t first, + std::size_t count, Layout layout) { + const auto local = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (local >= count) return; + size_t input_offset = first + local, output_offset = input_offset; + if (!layout.contiguous) { + size_t remaining = input_offset; + input_offset = output_offset = 0; + for (size_t axis = layout.rank; axis > 0; --axis) { + const auto coordinate = remaining % layout.shape[axis - 1]; + remaining /= layout.shape[axis - 1]; + input_offset += coordinate * layout.input_strides[axis - 1]; + output_offset += coordinate * layout.output_strides[axis - 1]; + } + } + out[output_offset] = + Caster::template Cast(input[input_offset]); +} +} // namespace hygon_cast_detail + +template <> +class Operator : public Cast { + public: + static bool supports(const Tensor input, Tensor out) { + if (input.device().type() != Device::Type::kHygon || + input.device() != out.device() || input.shape() != out.shape() || + !SupportedDtype(input.dtype()) || !SupportedDtype(out.dtype()) || + input.ndim() > 16 || !NonOverlapping(out)) + return false; + std::size_t input_bytes = 0, out_bytes = 0; + if (!SpanBytes(input, &input_bytes) || !SpanBytes(out, &out_bytes) || + !ValidAddress(input, input_bytes) || !ValidAddress(out, out_bytes)) + return false; + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess || + current != input.device().index()) + return false; + const auto ip = reinterpret_cast(input.data()); + const auto op = reinterpret_cast(out.data()); + const bool overlap = input_bytes && out_bytes && ip < op + out_bytes && + op < ip + input_bytes; + return !overlap || SameView(input, out); + } + + Operator(const Tensor input, Tensor out) + : Cast{Checked(input, out), out}, device_{input.device()} { + layout_.rank = input.ndim(); + layout_.contiguous = input.IsContiguous() && out.IsContiguous(); + for (size_t axis = 0; axis < input.ndim(); ++axis) { + layout_.shape[axis] = input.size(axis); + layout_.input_strides[axis] = input.stride(axis); + layout_.output_strides[axis] = out.stride(axis); + } + } + + void operator()(const Tensor input, Tensor out) const override { + if (input.dtype() != input_dtype_ || out.dtype() != out_dtype_ || + input.shape() != input_shape_ || out.shape() != out_shape_ || + input.strides() != input_strides_ || out.strides() != out_strides_ || + input.device() != device_ || !supports(input, out)) + throw std::invalid_argument( + "HYGON Cast: call metadata or alias changed after Make"); + if (output_size_ == 0 || SameView(input, out)) return; + const auto stream = static_cast(stream_); + DispatchFunc< + List, + List>( + {static_cast(input_dtype_), static_cast(out_dtype_)}, + [&](auto list_tag) { + using InT = TypeMapType(list_tag)>; + using OutT = TypeMapType(list_tag)>; + constexpr std::size_t cap = 65535u * 256u; + for (std::size_t first = 0; first < output_size_;) { + const auto count = std::min(cap, output_size_ - first); + const auto blocks = static_cast((count + 255) / 256); + hygon_cast_detail::Convert<<>>( + static_cast(input.data()), + static_cast(out.data()), first, count, layout_); + const auto status = cudaGetLastError(); + if (status != cudaSuccess) + throw std::runtime_error( + std::string("HYGON Cast: launch failed: ") + + cudaGetErrorString(status)); + first += count; + } + }, + "HYGON Cast"); + } + + private: + static bool SupportedDtype(DataType type) { + return type == DataType::kFloat32 || type == DataType::kFloat16 || + type == DataType::kBFloat16; + } + + static std::size_t Width(DataType type) { + return type == DataType::kFloat32 ? 4u : 2u; + } + + static bool SpanBytes(const Tensor t, std::size_t* bytes) { + if (t.shape().size() != t.strides().size()) return false; + bool empty = false; + for (std::size_t axis = 0; axis < t.ndim(); ++axis) { + if (t.stride(axis) < 0) return false; + empty = empty || t.size(axis) == 0; + } + if (empty) { + *bytes = 0; + return true; + } + const auto limit = + static_cast(std::numeric_limits::max()) / + Width(t.dtype()); + std::size_t span = 1; + for (std::size_t axis = 0; axis < t.ndim(); ++axis) { + const auto size = t.size(axis); + const auto stride = static_cast(t.stride(axis)); + if (size > 1 && stride > (limit - span) / (size - 1)) return false; + span += (size - 1) * stride; + } + *bytes = span * Width(t.dtype()); + return true; + } + + static bool NonOverlapping(const Tensor tensor) { + std::vector> axes; + for (size_t axis = 0; axis < tensor.ndim(); ++axis) { + if (tensor.stride(axis) < 0) return false; + if (tensor.size(axis) > 1) + axes.emplace_back(tensor.stride(axis), tensor.size(axis)); + } + std::sort(axes.begin(), axes.end()); + size_t extent = 1; + for (const auto& [stride, size] : axes) { + if (stride < extent || + (size - 1) > (std::numeric_limits::max() - extent) / stride) + return false; + extent += (size - 1) * stride; + } + return true; + } + + static bool ValidAddress(const Tensor t, std::size_t bytes) { + if (bytes == 0) return true; + const auto start = reinterpret_cast(t.data()); + return start && start % Width(t.dtype()) == 0 && + start <= std::numeric_limits::max() - bytes; + } + + static bool SameView(const Tensor input, const Tensor out) { + return input.dtype() == out.dtype() && input.data() == out.data() && + input.shape() == out.shape() && input.strides() == out.strides(); + } + + static Tensor Checked(const Tensor input, const Tensor out) { + if (!supports(input, out)) + throw std::invalid_argument( + "HYGON Cast: unsupported dtype/device/layout or alias"); + return input; + } + + Device device_; + hygon_cast_detail::Layout layout_; +}; +} // namespace infini::ops +#endif diff --git a/src/native/cuda/hygon/ops/causal_conv1d/kernel.h b/src/native/cuda/hygon/ops/causal_conv1d/kernel.h new file mode 100644 index 000000000..254d751e3 --- /dev/null +++ b/src/native/cuda/hygon/ops/causal_conv1d/kernel.h @@ -0,0 +1,202 @@ +#pragma once +#include +#include +#include + +#include + +#include "base/causal_conv1d.h" +#include "native/cuda/hygon/caster.cuh" +#include "native/cuda/hygon/runtime_.h" + +namespace infini::ops { +namespace causal_conv1d_detail { +struct Layout { + ptrdiff_t strides[9][3]{}; + size_t tokens, channels, pool; + bool cu_i64, initial_i64, final_i64; +}; +__device__ inline int64_t index(const void* p, bool i64, int i, int fallback) { + return p ? (i64 ? static_cast(p)[i] + : static_cast(static_cast(p)[i])) + : fallback; +} +template +__device__ inline float history(const T* state, const T* qkv, int64_t pos, + int64_t begin, int batch, int channel, + ptrdiff_t base, const Layout& p) { + if (pos < 3) return static_cast(state[base + pos * p.strides[2][2]]); + const auto off = static_cast(batch) * p.strides[0][0] + + static_cast(begin + pos - 3) * p.strides[0][1] + + static_cast(channel) * p.strides[0][2]; + return static_cast(qkv[off]); +} +// Port of baseline InfiniCore-0.2.9 causal_conv1d_k4_kernel. One thread per +// request/channel reads the entire sequence before committing its final state. +template +__global__ void kernel(const T* qkv, const T* weight, T* state, const T* bias, + const void* cu, const void* initial, + const void* final_indices, T* out, T* final_state, + Layout p) { + const int request = blockIdx.x, channel = blockIdx.y; + if (static_cast(channel) >= p.channels || threadIdx.x != 0) return; + int64_t begin = 0, end = static_cast(p.tokens); + int batch = request; + if (cu) { + begin = index(cu, p.cu_i64, request, 0); + end = index(cu, p.cu_i64, request + 1, 0); + batch = 0; + // Preserve baseline device-value behavior: malformed range leaves the + // corresponding output/state untouched. This is not host validation. + if (begin < 0 || end < begin || end > static_cast(p.tokens)) + return; + } + const int64_t length = end - begin; + const int64_t read = index(initial, p.initial_i64, request, request); + const int64_t write = index(final_indices, p.final_i64, request, request); + if (read < 0 || write < 0 || read >= static_cast(p.pool) || + (final_indices && write >= static_cast(p.pool))) + return; + const auto base = static_cast(read) * p.strides[2][0] + + static_cast(channel) * p.strides[2][1]; + const auto wb = static_cast(channel) * p.strides[1][0]; + for (int64_t t = 0; t < length; ++t) { + float acc = static_cast(weight[wb]) * + history(state, qkv, t, begin, batch, channel, base, p) + + static_cast(weight[wb + p.strides[1][2]]) * + history(state, qkv, t + 1, begin, batch, channel, base, p) + + static_cast(weight[wb + 2 * p.strides[1][2]]) * + history(state, qkv, t + 2, begin, batch, channel, base, p) + + static_cast(weight[wb + 3 * p.strides[1][2]]) * + history(state, qkv, t + 3, begin, batch, channel, base, p); + if (bias) + acc += static_cast( + bias[static_cast(channel) * p.strides[3][0]]); + const auto off = static_cast(batch) * p.strides[7][0] + + static_cast(begin + t) * p.strides[7][1] + + static_cast(channel) * p.strides[7][2]; + out[off] = static_cast(acc); + } + T* target = final_indices ? state : final_state; + const int stride_index = final_indices ? 2 : 8; + const auto target_base = + static_cast(final_indices ? write : request) * + p.strides[stride_index][0] + + static_cast(channel) * p.strides[stride_index][1]; + // Same ordered commit as baseline. Exact initial/final alias is valid for + // short sequences because copied history source positions lie ahead of + // writes. + target[target_base] = static_cast( + history(state, qkv, length, begin, batch, channel, base, p)); + target[target_base + p.strides[stride_index][2]] = static_cast( + history(state, qkv, length + 1, begin, batch, channel, base, p)); + target[target_base + 2 * p.strides[stride_index][2]] = static_cast( + history(state, qkv, length + 2, begin, batch, channel, base, p)); +} +} // namespace causal_conv1d_detail + +template <> +class Operator : public CausalConv1d { + public: + Operator(const Tensor qkv, const Tensor weight, Tensor conv_state, + const std::optional bias, + const std::optional cu_seqlens, + const std::optional initial_state_indices, + const std::optional final_state_indices, Tensor out, + std::optional final_conv_state) + : CausalConv1d(qkv, weight, conv_state, bias, cu_seqlens, + initial_state_indices, final_state_indices, out, + final_conv_state) { + if (device_.type() != Device::Type::kHygon) + throw std::invalid_argument("causal conv requires HYGON"); + if (weight.shape()[2] != 4 || conv_state.shape()[2] != 3) + throw std::invalid_argument( + "HYGON causal conv kernel requires K4 and state_len3"); + if (tokens_ > static_cast(std::numeric_limits::max() - 3)) + throw std::overflow_error( + "HYGON causal conv history index exceeds int64"); + check_device(); + // Baseline uses int request/channel variables and CUDA's grid.y dimension. + if (requests_ >= static_cast(INT_MAX) || channels_ > 65535) + throw std::invalid_argument( + "causal conv request/channel launch extent unsupported"); + p_.tokens = tokens_; + p_.channels = channels_; + p_.pool = pool_; + for (size_t i = 0; i < metadata_.size(); ++i) + if (metadata_[i]) + for (size_t j = 0; j < metadata_[i]->strides().size(); ++j) + p_.strides[i][j] = metadata_[i]->strides()[j]; + p_.cu_i64 = metadata_[4] && metadata_[4]->dtype() == DataType::kInt64; + p_.initial_i64 = metadata_[5] && metadata_[5]->dtype() == DataType::kInt64; + p_.final_i64 = metadata_[6] && metadata_[6]->dtype() == DataType::kInt64; + } + // Device capture/state-pool lifetime is not supported for this stateful path. + // Zero device allocations does not alone establish capture eligibility. + void operator()(const Tensor qkv, const Tensor weight, Tensor conv_state, + const std::optional bias, + const std::optional cu_seqlens, + const std::optional initial_state_indices, + const std::optional final_state_indices, Tensor out, + std::optional final_conv_state) const override { + const Views views{qkv, + weight, + conv_state, + bias, + cu_seqlens, + initial_state_indices, + final_state_indices, + out, + final_conv_state}; + for (size_t i = 0; i < views.size(); ++i) { + const auto& t = views[i]; + if (t.has_value() != metadata_[i].has_value() || + (t && (!std::equal_to{}(*t, *metadata_[i]) || + (t->numel() && !t->data())))) + throw std::invalid_argument( + "causal conv metadata changed or nonempty storage is null"); + } + check_device(); + if (!requests_ || !channels_) return; + switch (dtype_) { + case DataType::kFloat32: + launch(views); + break; + case DataType::kFloat16: + launch(views); + break; + case DataType::kBFloat16: + launch<__nv_bfloat16>(views); + break; + default: + throw std::invalid_argument("causal conv dtype unsupported"); + } + if (cudaPeekAtLastError() != cudaSuccess) + throw std::runtime_error("causal conv CUDA launch failed"); + } + + private: + causal_conv1d_detail::Layout p_{}; + void check_device() const { + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess) + throw std::runtime_error("causal conv device query failed"); + if (current != device_.index()) + throw std::invalid_argument("causal conv owner device mismatch"); + } + template + void launch(const Views& v) const { + auto pointer = [&](size_t i) -> void* { + return v[i] ? const_cast(v[i]->data()) : nullptr; + }; + const dim3 grid(static_cast(requests_), + static_cast(channels_), 1); + const auto stream = static_cast(stream_ ? stream_ : nullptr); + causal_conv1d_detail::kernel<<>>( + static_cast(pointer(0)), static_cast(pointer(1)), + static_cast(pointer(2)), static_cast(pointer(3)), + pointer(4), pointer(5), pointer(6), static_cast(pointer(7)), + static_cast(pointer(8)), p_); + } +}; +} // namespace infini::ops diff --git a/src/native/cuda/hygon/ops/fill/kernel.h b/src/native/cuda/hygon/ops/fill/kernel.h new file mode 100644 index 000000000..09a9ec9c6 --- /dev/null +++ b/src/native/cuda/hygon/ops/fill/kernel.h @@ -0,0 +1,16 @@ +#pragma once + +#include "native/cuda/hygon/caster.cuh" +#include "native/cuda/hygon/runtime_.h" +#include "native/cuda/ops/fill/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaFill> { + public: + using CudaFill>::CudaFill; +}; + +} // namespace infini::ops diff --git a/src/native/cuda/hygon/ops/fused_gated_delta_net_gating/kernel.h b/src/native/cuda/hygon/ops/fused_gated_delta_net_gating/kernel.h new file mode 100644 index 000000000..8b35c724d --- /dev/null +++ b/src/native/cuda/hygon/ops/fused_gated_delta_net_gating/kernel.h @@ -0,0 +1,120 @@ +#pragma once +#include +#include +#include + +#include +#include +#include + +#include "base/fused_gated_delta_net_gating.h" +#include "native/cuda/hygon/caster.cuh" +#include "native/cuda/hygon/runtime_.h" + +namespace infini::ops { +namespace gating_detail { +struct Layout { + ptrdiff_t strides[6][3]{}; + size_t total, seq, hidden; + float beta, threshold; +}; +template +__global__ void kernel(const T* A, const T* a, const T* b, const T* dt, + float* g, float* betao, Layout p) { + const size_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (linear >= p.total) return; + const size_t h = linear % p.hidden, s = (linear / p.hidden) % p.seq; + const size_t batch = linear / p.hidden / p.seq; + ptrdiff_t offsets[6]; + for (int i = 0; i < 6; ++i) + offsets[i] = static_cast(batch) * p.strides[i][0] + + static_cast(s) * p.strides[i][1] + + static_cast(h) * p.strides[i][2]; + const float x = static_cast(a[offsets[1]]) + + static_cast(dt[h * p.strides[3][0]]); + const float bx = p.beta * x; + const float softplus = bx <= p.threshold ? log1pf(expf(bx)) / p.beta : x; + g[offsets[4]] = -expf(static_cast(A[h * p.strides[0][0]])) * softplus; + const float bv = static_cast(b[offsets[2]]); + const float z = expf(bv >= 0 ? -bv : bv); + betao[offsets[5]] = bv >= 0 ? 1.0f / (1.0f + z) : z / (1.0f + z); +} +} // namespace gating_detail + +template <> +class Operator + : public FusedGatedDeltaNetGating { + public: + Operator(Tensor A, Tensor a, Tensor b, Tensor dt, float beta, float threshold, + Tensor g, Tensor betao) + : FusedGatedDeltaNetGating(A, a, b, dt, beta, threshold, g, betao) { + if (device_.type() != Device::Type::kHygon) + throw std::invalid_argument("gating requires HYGON device"); + check_device(); + if (total_ / 256 + (total_ % 256 != 0) > + std::numeric_limits::max()) + throw std::overflow_error("gating grid exceeds CUDA limit"); + p_.total = total_; + p_.seq = seq_; + p_.hidden = hidden_; + p_.beta = beta; + p_.threshold = threshold; + for (size_t i = 0; i < 6; ++i) + for (size_t j = 0; j < metadata_[i].strides().size(); ++j) + p_.strides[i][j] = metadata_[i].strides()[j]; + } + // All metadata is a kernel value argument. There is no allocation, host copy, + // synchronization or data-dependent host work in operator(); stream launch + // is capturable. Make owns no device resources and requires zero workspace. + bool capture_safe() const override { return true; } + void operator()(Tensor A, Tensor a, Tensor b, Tensor dt, float beta, + float threshold, Tensor g, Tensor betao) const override { + check_device(); + std::array tensors{A, a, b, dt, g, betao}; + for (size_t i = 0; i < 6; ++i) + if (!std::equal_to{}(tensors[i], metadata_[i])) + throw std::invalid_argument("gating metadata changed after Make"); + if (std::memcmp(&beta, &beta_, sizeof(float)) || + std::memcmp(&threshold, &threshold_, sizeof(float))) + throw std::invalid_argument("gating attributes changed after Make"); + if (!total_) return; + const dim3 grid( + static_cast(total_ / 256 + (total_ % 256 != 0))); + const auto stream = static_cast(stream_ ? stream_ : nullptr); + switch (metadata_[1].dtype()) { + case DataType::kFloat32: + launch(A, a, b, dt, g, betao, grid, stream); + break; + case DataType::kFloat16: + launch(A, a, b, dt, g, betao, grid, stream); + break; + case DataType::kBFloat16: + launch<__nv_bfloat16>(A, a, b, dt, g, betao, grid, stream); + break; + default: + throw std::invalid_argument("gating dtype unsupported"); + } + if (cudaPeekAtLastError() != cudaSuccess) + throw std::runtime_error("gating CUDA launch failed"); + } + + private: + gating_detail::Layout p_{}; + void check_device() const { + int current = -1; + if (cudaGetDevice(¤t) != cudaSuccess) + throw std::runtime_error("gating device query failed"); + if (current != device_.index()) + throw std::invalid_argument("gating owner device mismatch"); + } + template + void launch(Tensor A, Tensor a, Tensor b, Tensor dt, Tensor g, Tensor betao, + dim3 grid, cudaStream_t stream) const { + gating_detail::kernel<<>>( + static_cast(A.data()), static_cast(a.data()), + static_cast(b.data()), static_cast(dt.data()), + static_cast(g.data()), static_cast(betao.data()), p_); + } +}; +} // namespace infini::ops diff --git a/src/native/cuda/hygon/ops/moe_align_block_size/kernel.cu b/src/native/cuda/hygon/ops/moe_align_block_size/kernel.cu new file mode 100644 index 000000000..76f5938a9 --- /dev/null +++ b/src/native/cuda/hygon/ops/moe_align_block_size/kernel.cu @@ -0,0 +1,81 @@ +#include + +#include +#include +#include + +#include "native/cuda/hygon/ops/moe_align_block_size/kernel.cuh" +#include "native/cuda/hygon/ops/moe_align_block_size/kernel.h" + +namespace infini::ops { +namespace { + +class DeviceGuard { + public: + explicit DeviceGuard(int device_index) { + auto status = cudaGetDevice(&previous_device_); + assert(status == cudaSuccess && + "`MoeAlignBlockSize` failed to query the current CUDA device"); + + if (previous_device_ != device_index) { + status = cudaSetDevice(device_index); + assert(status == cudaSuccess && + "`MoeAlignBlockSize` failed to select the input CUDA device"); + restore_ = true; + } + } + + ~DeviceGuard() { + if (restore_) { + const auto status = cudaSetDevice(previous_device_); + assert(status == cudaSuccess && + "`MoeAlignBlockSize` failed to restore the CUDA device"); + } + } + + private: + int previous_device_{0}; + + bool restore_{false}; +}; + +} // namespace + +Operator::Operator( + const Tensor topk_ids, const std::optional expert_map, + const int64_t num_experts, const int64_t block_size, + Tensor sorted_token_ids, Tensor experts_ids, Tensor num_tokens_post_pad) + : MoeAlignBlockSize{topk_ids, expert_map, num_experts, + block_size, sorted_token_ids, experts_ids, + num_tokens_post_pad}, + device_index_{topk_ids.device().index()} {} + +void Operator::Run( + const Tensor topk_ids, const std::optional maybe_expert_map, + const int64_t num_experts, const int64_t block_size, + Tensor sorted_token_ids, Tensor experts_ids, + Tensor num_tokens_post_pad) const { + DeviceGuard device_guard{device_index_}; + + constexpr int32_t kThreads = 256; + const auto shared_memory_size = + static_cast(num_experts_) * sizeof(int32_t); + moe_align_block_size_detail::MoeAlignBlockSizeKernel<<< + 1, kThreads, shared_memory_size, static_cast(stream_)>>>( + reinterpret_cast(topk_ids.data()), + maybe_expert_map + ? reinterpret_cast(maybe_expert_map->data()) + : nullptr, + reinterpret_cast(sorted_token_ids.data()), + reinterpret_cast(experts_ids.data()), + reinterpret_cast(num_tokens_post_pad.data()), numel_, + static_cast(num_experts_), static_cast(block_size_), + sorted_token_ids_size_, experts_ids_size_, + maybe_expert_map ? static_cast(maybe_expert_map->numel()) : 0); + + const auto status = cudaGetLastError(); + assert(status == cudaSuccess && + "`MoeAlignBlockSize` CUDA kernel launch failed"); +} + +} // namespace infini::ops diff --git a/src/native/cuda/hygon/ops/moe_align_block_size/kernel.cuh b/src/native/cuda/hygon/ops/moe_align_block_size/kernel.cuh new file mode 100644 index 000000000..2055342e1 --- /dev/null +++ b/src/native/cuda/hygon/ops/moe_align_block_size/kernel.cuh @@ -0,0 +1,86 @@ +#ifndef INFINI_OPS_HYGON_MOE_ALIGN_BLOCK_SIZE_KERNEL_CUH_ +#define INFINI_OPS_HYGON_MOE_ALIGN_BLOCK_SIZE_KERNEL_CUH_ + +#include +#include + +namespace infini::ops { +namespace moe_align_block_size_detail { + +__device__ __forceinline__ int32_t MapExpertId(int32_t expert_id, + const int32_t* expert_map, + int32_t num_experts, + int32_t map_size) { + if (expert_id < 0 || expert_id >= (expert_map ? map_size : num_experts)) { + return -1; + } + if (expert_map != nullptr) { + expert_id = expert_map[expert_id]; + } + + return expert_id >= 0 && expert_id < num_experts ? expert_id : -1; +} + +__global__ void MoeAlignBlockSizeKernel( + const int32_t* __restrict__ topk_ids, + const int32_t* __restrict__ expert_map, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ experts_ids, + int32_t* __restrict__ num_tokens_post_pad, size_t numel, + int32_t num_experts, int32_t block_size, size_t sorted_token_ids_size, + size_t experts_ids_size, int32_t map_size) { + extern __shared__ int32_t expert_offsets[]; + + for (size_t i = threadIdx.x; i < sorted_token_ids_size; i += blockDim.x) { + sorted_token_ids[i] = static_cast(numel); + } + for (size_t i = threadIdx.x; i < experts_ids_size; i += blockDim.x) { + experts_ids[i] = -1; + } + for (int32_t i = threadIdx.x; i < num_experts; i += blockDim.x) { + expert_offsets[i] = 0; + } + if (threadIdx.x == 0) { + *num_tokens_post_pad = 0; + } + __syncthreads(); + + for (size_t i = threadIdx.x; i < numel; i += blockDim.x) { + const auto expert_id = + MapExpertId(topk_ids[i], expert_map, num_experts, map_size); + if (expert_id >= 0) { + atomicAdd(&expert_offsets[expert_id], 1); + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int32_t token_offset = 0; + for (int32_t expert_id = 0; expert_id < num_experts; ++expert_id) { + const int32_t count = expert_offsets[expert_id]; + expert_offsets[expert_id] = token_offset; + const int32_t padded_count = + (count + block_size - 1) / block_size * block_size; + for (int32_t i = 0; i < padded_count; i += block_size) { + experts_ids[(token_offset + i) / block_size] = expert_id; + } + token_offset += padded_count; + } + *num_tokens_post_pad = token_offset; + } + __syncthreads(); + + if (threadIdx.x == 0) { + for (size_t i = 0; i < numel; ++i) { + const auto expert_id = + MapExpertId(topk_ids[i], expert_map, num_experts, map_size); + if (expert_id >= 0) { + sorted_token_ids[expert_offsets[expert_id]++] = static_cast(i); + } + } + } +} + +} // namespace moe_align_block_size_detail +} // namespace infini::ops + +#endif // INFINI_OPS_HYGON_MOE_ALIGN_BLOCK_SIZE_KERNEL_CUH_ diff --git a/src/native/cuda/hygon/ops/moe_align_block_size/kernel.h b/src/native/cuda/hygon/ops/moe_align_block_size/kernel.h new file mode 100644 index 000000000..0984b5205 --- /dev/null +++ b/src/native/cuda/hygon/ops/moe_align_block_size/kernel.h @@ -0,0 +1,35 @@ +#ifndef INFINI_OPS_HYGON_MOE_ALIGN_BLOCK_SIZE_KERNEL_H_ +#define INFINI_OPS_HYGON_MOE_ALIGN_BLOCK_SIZE_KERNEL_H_ + +#include + +#include "base/moe_align_block_size.h" + +namespace infini::ops { + +template <> +class Operator + : public MoeAlignBlockSize { + public: + Operator(const Tensor topk_ids, const std::optional expert_map, + const int64_t num_experts, const int64_t block_size, + Tensor sorted_token_ids, Tensor experts_ids, + Tensor num_tokens_post_pad); + + using MoeAlignBlockSize::operator(); + + bool capture_safe() const override { return true; } + + protected: + void Run(const Tensor topk_ids, const std::optional maybe_expert_map, + const int64_t num_experts, const int64_t block_size, + Tensor sorted_token_ids, Tensor experts_ids, + Tensor num_tokens_post_pad) const override; + + private: + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_HYGON_MOE_ALIGN_BLOCK_SIZE_KERNEL_H_ diff --git a/src/native/cuda/hygon/ops/mrope/kernel.h b/src/native/cuda/hygon/ops/mrope/kernel.h new file mode 100644 index 000000000..7808c0627 --- /dev/null +++ b/src/native/cuda/hygon/ops/mrope/kernel.h @@ -0,0 +1,12 @@ +#pragma once +#include "native/cuda/hygon/caster.cuh" +#include "native/cuda/hygon/runtime_.h" +#include "native/cuda/ops/mrope/kernel.h" +namespace infini::ops { +template <> +class Operator + : public CudaMrope> { + public: + using CudaMrope>::CudaMrope; +}; +} // namespace infini::ops diff --git a/src/native/cuda/hygon/ops/mul/kernel.h b/src/native/cuda/hygon/ops/mul/kernel.h new file mode 100644 index 000000000..b949617a7 --- /dev/null +++ b/src/native/cuda/hygon/ops/mul/kernel.h @@ -0,0 +1,16 @@ +#pragma once + +#include "native/cuda/hygon/caster.cuh" +#include "native/cuda/hygon/runtime_.h" +#include "native/cuda/ops/mul/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaMul> { + public: + using CudaMul>::CudaMul; +}; + +} // namespace infini::ops diff --git a/src/native/cuda/hygon/ops/separate_table_rope/kernel.h b/src/native/cuda/hygon/ops/separate_table_rope/kernel.h new file mode 100644 index 000000000..d4aa1316e --- /dev/null +++ b/src/native/cuda/hygon/ops/separate_table_rope/kernel.h @@ -0,0 +1,27 @@ +#ifndef INFINI_OPS_HYGON_SEPARATE_TABLE_ROPE_KERNEL_H_ +#define INFINI_OPS_HYGON_SEPARATE_TABLE_ROPE_KERNEL_H_ +#include "native/cuda/hygon/caster.cuh" +#include "native/cuda/hygon/runtime_.h" +#include "native/cuda/ops/separate_table_rope/kernel.h" +namespace infini::ops { +template <> +class Operator + : public CudaSeparateTableRope> { + public: + using CudaSeparateTableRope< + Runtime>::CudaSeparateTableRope; + void operator()(const Tensor x, const Tensor p, const Tensor s, + const Tensor c, bool neox, Tensor y) const override { + CudaSeparateTableRope>::operator()(x, p, s, c, + neox, y); + if (x.numel() == 0) return; + // Same native HYGON launch-error API used by the installed Argmax source. + // Peek reports immediate launch failure; it does not synchronize or claim + // that device data access / position bounds are valid. + const auto status = cudaPeekAtLastError(); + if (status != cudaSuccess) + throw std::runtime_error(cudaGetErrorString(status)); + } +}; +} // namespace infini::ops +#endif diff --git a/src/native/cuda/hygon/ops/tril/kernel.h b/src/native/cuda/hygon/ops/tril/kernel.h new file mode 100644 index 000000000..532ccd490 --- /dev/null +++ b/src/native/cuda/hygon/ops/tril/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_HYGON_TRIL_KERNEL_H_ +#define INFINI_OPS_HYGON_TRIL_KERNEL_H_ + +#include "native/cuda/hygon/runtime_.h" +#include "native/cuda/ops/triangular/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public triangular_detail::CudaTriangular, + Tril, false> { + public: + using triangular_detail::CudaTriangular, Tril, + false>::CudaTriangular; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_HYGON_TRIL_KERNEL_H_ diff --git a/src/native/cuda/hygon/ops/triu/kernel.h b/src/native/cuda/hygon/ops/triu/kernel.h new file mode 100644 index 000000000..aa237ba7a --- /dev/null +++ b/src/native/cuda/hygon/ops/triu/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_HYGON_TRIU_KERNEL_H_ +#define INFINI_OPS_HYGON_TRIU_KERNEL_H_ + +#include "native/cuda/hygon/runtime_.h" +#include "native/cuda/ops/triangular/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public triangular_detail::CudaTriangular, + Triu, true> { + public: + using triangular_detail::CudaTriangular, Triu, + true>::CudaTriangular; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_HYGON_TRIU_KERNEL_H_ diff --git a/src/triton/jit/backend.h b/src/triton/jit/backend.h index ee080734f..21cf6ee62 100644 --- a/src/triton/jit/backend.h +++ b/src/triton/jit/backend.h @@ -31,6 +31,8 @@ struct Target { int warp_size{0}; }; +Target HygonTarget(int device_id); + enum class DataType { kPointer, kInt8, @@ -111,6 +113,11 @@ struct BackendTraits { static constexpr const char* kName = "cuda"; }; +template <> +struct BackendTraits { + static constexpr const char* kName = "hip"; +}; + template ::kName)>> class Backend { @@ -124,6 +131,9 @@ class Backend { static Target CurrentTarget() { const int device_id = CurrentDevice(); if (device_id < 0) return {}; + if constexpr (kDev == Device::Type::kHygon) { + return HygonTarget(device_id); + } int major = 0; auto status = Runtime::DeviceGetAttribute( @@ -179,12 +189,12 @@ class Backend { constexpr unsigned kDefaultSharedMemoryLimit = 48 * 1024; if (metadata.shared_memory_size > kDefaultSharedMemoryLimit) { int maximum_shared_memory = 0; - status = Runtime::DeviceGetAttribute( + const auto runtime_status = Runtime::DeviceGetAttribute( &maximum_shared_memory, Runtime::kDevAttrMaxSharedMemoryPerBlockOptin, CurrentDevice()); - assert(status == Runtime::kSuccess && + assert(runtime_status == Runtime::kSuccess && "Triton JIT failed to get the shared memory limit."); - if (status != Runtime::kSuccess) { + if (runtime_status != Runtime::kSuccess) { Driver::ModuleUnload(module); return nullptr; } diff --git a/src/triton/jit/compile.py b/src/triton/jit/compile.py index 8b49f8e94..6439d9bce 100644 --- a/src/triton/jit/compile.py +++ b/src/triton/jit/compile.py @@ -14,9 +14,28 @@ _DRIVER_BACKENDS = { "cuda": "nvidia", + "hip": "amd", } +def get_hygon_target(device_id): + """Read the compiler target from the active DTK Triton driver.""" + driver = triton.runtime.driver.active + previous = driver.get_current_device() + try: + driver.set_current_device(device_id) + target = driver.get_current_target() + if target.backend != "hip": + raise RuntimeError("Hygon JIT requires a HIP Triton target.") + return { + "backend": target.backend, + "architecture": target.arch, + "warp_size": target.warp_size, + } + finally: + driver.set_current_device(previous) + + def _cache_directory(): configured_path = os.environ.get("INFINI_OPS_TRITON_CACHE_DIR") if configured_path: @@ -211,7 +230,13 @@ def _use_device(target, device_id, stream): previous_driver = triton.runtime.driver.active previous_device = previous_driver.get_current_device() driver_backend = _DRIVER_BACKENDS[target["backend"]] - active_driver = triton.backends.backends[driver_backend].driver() + if target["backend"] == "hip": + # DTK supplies a vendor driver with its own device-property ABI. + active_driver = previous_driver + if active_driver.get_current_target().backend != "hip": + raise RuntimeError("HIP auto-tuning requires the active HIP driver.") + else: + active_driver = triton.backends.backends[driver_backend].driver() try: triton.runtime.driver.set_active(active_driver) @@ -312,8 +337,12 @@ def _write_artifacts(compiled_kernel, target, output_prefix, cache_identity): "num_warps": compiled_kernel.metadata.num_warps, "binary_ext": binary_extension, "shared": compiled_kernel.metadata.shared, - "global_scratch_size": compiled_kernel.metadata.global_scratch_size, - "profile_scratch_size": compiled_kernel.metadata.profile_scratch_size, + "global_scratch_size": getattr( + compiled_kernel.metadata, "global_scratch_size", 0 + ), + "profile_scratch_size": getattr( + compiled_kernel.metadata, "profile_scratch_size", 0 + ), "cache_identity": cache_identity, } temporary_paths = [] diff --git a/src/triton/jit/compiler.cc b/src/triton/jit/compiler.cc index d22323661..966a4a788 100644 --- a/src/triton/jit/compiler.cc +++ b/src/triton/jit/compiler.cc @@ -142,6 +142,23 @@ py::list BuildArguments(const std::vector& arguments) { } // namespace +Target HygonTarget(int device_id) { + static thread_local std::unordered_map targets; + const auto cached = targets.find(device_id); + if (cached != targets.end()) return cached->second; + const auto target = RunPython("Hygon target lookup", [&] { + const auto value = CompilerFunction("get_hygon_target")(device_id); + return Target{value["backend"].cast(), + value["architecture"].cast(), + value["warp_size"].cast()}; + }); + assert(target.has_value() && + "Triton JIT failed to identify the Hygon target."); + if (!target) return {}; + targets.emplace(device_id, *target); + return *target; +} + std::string CompilationFingerprint(const std::string& op_name) { static std::mutex cache_mutex; static std::unordered_map cache; diff --git a/src/triton/ops/add/jit.h b/src/triton/ops/add/jit.h index da3274c5c..821dbe1fb 100644 --- a/src/triton/ops/add/jit.h +++ b/src/triton/ops/add/jit.h @@ -95,6 +95,13 @@ class Operator using triton::jit::Add::Add; }; +template <> +class Operator + : public triton::jit::Add { + public: + using triton::jit::Add::Add; +}; + } // namespace infini::ops #endif diff --git a/src/triton/ops/clover_block_conv1d/jit.h b/src/triton/ops/clover_block_conv1d/jit.h new file mode 100644 index 000000000..d37c1d01d --- /dev/null +++ b/src/triton/ops/clover_block_conv1d/jit.h @@ -0,0 +1,56 @@ +#pragma once + +#include "base/clover_block_conv1d.h" +#include "triton/jit/jit.h" + +namespace infini::ops::triton::jit { + +template +class CloverBlockConv1d : public ::infini::ops::CloverBlockConv1d { + public: + using ::infini::ops::CloverBlockConv1d::CloverBlockConv1d; + + void operator()(const Tensor input, const Tensor state, const Tensor indices, + const Tensor weight, const Tensor right_weight, + const Tensor valid, size_t block_size, + Tensor out) const override { + if (input.numel() == 0) return; + const Config config{4u, + 1u, + {{"NUM_CHANNELS", input.size(1)}, + {"BLOCK_SIZE", block_size}, + {"NUM_SLOTS", state.size(0)}, + {"WIDTH", 4}, + {"HAS_BIAS", 0}, + {"INDEPENDENT_SYMMETRIC", 1}, + {"BLOCK_T", 4}, + {"NUM_TOKEN_TILES", (block_size + 3) / 4}, + {"BLOCK_C", 128}}}; + Launch( + "clover_block_conv1d", input.device().index(), this->stream_, + this->config_ptr_.get(), config, {}, + [&](const Config&) -> std::optional { + return Grid{ + static_cast((input.size(1) + 127) / 128), + static_cast(indices.numel() * ((block_size + 3) / 4))}; + }, + input, state, indices, weight, weight, right_weight, valid, out, + input.stride(0), input.stride(1), state.stride(0), state.stride(1), + state.stride(2), weight.stride(0), weight.stride(2), + right_weight.stride(0), right_weight.stride(2), out.stride(0), + out.stride(1)); + } +}; + +} // namespace infini::ops::triton::jit + +namespace infini::ops { + +template <> +class Operator + : public triton::jit::CloverBlockConv1d { + public: + using triton::jit::CloverBlockConv1d::CloverBlockConv1d; +}; + +} // namespace infini::ops diff --git a/src/triton/ops/clover_block_conv1d/jit.py b/src/triton/ops/clover_block_conv1d/jit.py new file mode 100644 index 000000000..898fa496b --- /dev/null +++ b/src/triton/ops/clover_block_conv1d/jit.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from SGLang PR #46927, commit 058939af24d5cdc44bb4b8852d0ddfab555ed71e. +# Changes: InfiniOps JIT entry point, Hygon launch geometry and pool bounds. +import triton +import triton.language as tl + + +@triton.jit +def clover_block_conv1d_kernel( + input_ptr, + clean_state_ptr, + cache_indices_ptr, + weight_ptr, + bias_ptr, + right_weight_ptr, + active_valid_ptr, + output_ptr, + input_stride_t, + input_stride_c, + state_stride_slot, + state_stride_c, + state_stride_w, + weight_stride_c, + weight_stride_w, + right_weight_stride_c, + right_weight_stride_w, + output_stride_t, + output_stride_c, + NUM_CHANNELS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + NUM_SLOTS: tl.constexpr, + WIDTH: tl.constexpr, + HAS_BIAS: tl.constexpr, + INDEPENDENT_SYMMETRIC: tl.constexpr, + BLOCK_T: tl.constexpr, + NUM_TOKEN_TILES: tl.constexpr, + BLOCK_C: tl.constexpr, +): + channel_offsets = tl.program_id(0) * BLOCK_C + tl.arange(0, BLOCK_C) + token_program = tl.program_id(1) + request_idx = token_program // NUM_TOKEN_TILES + token_tile = token_program % NUM_TOKEN_TILES + token_offsets = token_tile * BLOCK_T + tl.arange(0, BLOCK_T) + channel_mask = channel_offsets < NUM_CHANNELS + token_mask = token_offsets < BLOCK_SIZE + token_indices = request_idx * BLOCK_SIZE + token_offsets + + cache_index = tl.load(cache_indices_ptr + request_idx) + request_active = (cache_index >= 0) & (cache_index < NUM_SLOTS) + safe_cache_index = tl.maximum(cache_index, 0) + output_active = ( + tl.load(active_valid_ptr + token_indices, mask=token_mask, other=0) != 0 + ) & request_active + + if HAS_BIAS: + bias = tl.load(bias_ptr + channel_offsets, mask=channel_mask, other=0.0).to( + tl.float32 + ) + else: + bias = tl.zeros((BLOCK_C,), dtype=tl.float32) + accumulator = tl.broadcast_to(bias[None, :], (BLOCK_T, BLOCK_C)) + + for tap in tl.static_range(WIDTH): + relative_source = token_offsets + tap - (WIDTH - 1) + from_current = relative_source >= 0 + current_token = request_idx * BLOCK_SIZE + tl.maximum(relative_source, 0) + current_valid = ( + tl.load( + active_valid_ptr + current_token, + mask=token_mask & from_current, + other=0, + ) + != 0 + ) & request_active + current_value = tl.load( + input_ptr + + current_token[:, None] * input_stride_t + + channel_offsets[None, :] * input_stride_c, + mask=(token_mask & from_current & current_valid)[:, None] + & channel_mask[None, :], + other=0.0, + ).to(tl.float32) + + history_offset = relative_source + (WIDTH - 1) + history_value = tl.load( + clean_state_ptr + + safe_cache_index * state_stride_slot + + channel_offsets[None, :] * state_stride_c + + tl.maximum(history_offset, 0)[:, None] * state_stride_w, + mask=(token_mask & ~from_current & request_active)[:, None] + & channel_mask[None, :], + other=0.0, + ).to(tl.float32) + source_value = tl.where(from_current[:, None], current_value, history_value) + weight = tl.load( + weight_ptr + channel_offsets * weight_stride_c + tap * weight_stride_w, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + accumulator += source_value * weight[None, :] + + if INDEPENDENT_SYMMETRIC: + right_accumulator = tl.zeros((BLOCK_T, BLOCK_C), dtype=tl.float32) + for tap in tl.static_range(WIDTH - 1): + distance = WIDTH - 1 - tap + right_offset = token_offsets + distance + right_token = request_idx * BLOCK_SIZE + right_offset + right_valid = ( + tl.load( + active_valid_ptr + right_token, + mask=token_mask & (right_offset < BLOCK_SIZE), + other=0, + ) + != 0 + ) & request_active + right_value = tl.load( + input_ptr + + right_token[:, None] * input_stride_t + + channel_offsets[None, :] * input_stride_c, + mask=(token_mask & (right_offset < BLOCK_SIZE) & right_valid)[:, None] + & channel_mask[None, :], + other=0.0, + ).to(tl.float32) + weight = tl.load( + right_weight_ptr + + channel_offsets * right_weight_stride_c + + (distance - 1) * right_weight_stride_w, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + right_accumulator += right_value * weight[None, :] + accumulator += right_accumulator + + output = accumulator * tl.sigmoid(accumulator) + output = tl.where(output_active[:, None], output, 0.0) + tl.store( + output_ptr + + token_indices[:, None] * output_stride_t + + channel_offsets[None, :] * output_stride_c, + output, + mask=token_mask[:, None] & channel_mask[None, :], + ) + + +kernel = clover_block_conv1d_kernel diff --git a/src/triton/ops/clover_coreset/jit.h b/src/triton/ops/clover_coreset/jit.h new file mode 100644 index 000000000..bf000bd3f --- /dev/null +++ b/src/triton/ops/clover_coreset/jit.h @@ -0,0 +1,46 @@ +#pragma once + +#include "base/clover_coreset.h" +#include "triton/jit/jit.h" + +namespace infini::ops { + +template <> +class Operator : public CloverCoreset { + public: + Operator(const Tensor logits, size_t block_size, double top_p, Tensor out) + : CloverCoreset(logits, block_size, top_p, out), + blocks_(logits.size(0) / block_size) {} + + size_t workspace_size_in_bytes() const override { return blocks_ * 256 * 5; } + + void operator()(const Tensor logits, size_t block_size, double top_p, + Tensor out) const override { + if (!blocks_) return; + if (!workspace_ || workspace_size_in_bytes_ < workspace_size_in_bytes()) + throw std::invalid_argument("CloverCoreset: missing workspace."); + Tensor mass{ + workspace_, {blocks_, 256}, DataType::kFloat32, logits.device()}; + Tensor allow{static_cast(workspace_) + blocks_ * 256 * 4, + {blocks_, 256}, + DataType::kInt8, + logits.device()}; + for (int stage = 0; stage < 3; ++stage) { + const triton::jit::Config config{ + 4u, 1u, {{"WIDTH", block_size}, {"STAGE", stage}}}; + triton::jit::Launch( + "clover_coreset", logits.device().index(), stream_, nullptr, config, + {}, + [&](const triton::jit::Config&) -> std::optional { + return triton::jit::Grid{ + static_cast(stage == 2 ? logits.size(0) : blocks_)}; + }, + logits, mass, allow, out, static_cast(top_p)); + } + } + + private: + size_t blocks_; +}; + +} // namespace infini::ops diff --git a/src/triton/ops/clover_coreset/jit.py b/src/triton/ops/clover_coreset/jit.py new file mode 100644 index 000000000..08dc7d469 --- /dev/null +++ b/src/triton/ops/clover_coreset/jit.py @@ -0,0 +1,52 @@ +import triton +import triton.language as tl + + +@triton.jit +def clover_coreset_kernel( + Logits, Mass, Allow, Out, top_p, WIDTH: tl.constexpr, STAGE: tl.constexpr +): + expert = tl.arange(0, 256) + if STAGE == 0: + block = tl.program_id(0) + tokens = tl.arange(0, 32) + logits = tl.load( + Logits + (block * WIDTH + tokens[:, None]) * 256 + expert[None, :], + tokens[:, None] < WIDTH, + -float("inf"), + ).to(tl.float32) + row_max = tl.max(logits, 1) + probability = tl.exp(logits - row_max[:, None]) + probability = probability / tl.sum(probability, 1)[:, None] + selected = tl.full((32, 256), False, tl.int1) + remaining = logits + for rank in tl.static_range(8): + best = tl.argmax(remaining, 1, tie_break_left=True) + chosen = expert[None, :] == best[:, None] + selected |= chosen + remaining = tl.where(chosen, -float("inf"), remaining) + mass = tl.sum( + tl.where(selected & (tokens[:, None] < WIDTH), probability, 0.0), 0 + ) + tl.store(Mass + block * 256 + expert, mass) + elif STAGE == 1: + block = tl.program_id(0) + mass = tl.load(Mass + block * 256 + expert) + # Positive float bits sort by magnitude; low bits break ties by expert ID. + bits = mass.to(tl.int32, bitcast=True).to(tl.int64) + key = ((0x7FFFFFFF - bits) << 32) | expert.to(tl.int64) + order = (tl.sort(key, descending=False) & 0xFFFFFFFF).to(tl.int32) + sorted_mass = tl.load(Mass + block * 256 + order) + probabilities = sorted_mass / tl.maximum(tl.sum(mass, 0), 1.0e-9) + keep = tl.cumsum(probabilities, 0) - probabilities < top_p + keep = tl.where(top_p >= 1.0, sorted_mass > 0.0, keep) + keep = keep | (expert == 0) + tl.store(Allow + block * 256 + order, keep) + else: + token = tl.program_id(0) + logits = tl.load(Logits + token * 256 + expert) + keep = tl.load(Allow + (token // WIDTH) * 256 + expert) + tl.store(Out + token * 256 + expert, tl.where(keep != 0, logits, -float("inf"))) + + +kernel = clover_coreset_kernel diff --git a/src/triton/ops/clover_live_mask/jit.h b/src/triton/ops/clover_live_mask/jit.h new file mode 100644 index 000000000..c8b0a65c3 --- /dev/null +++ b/src/triton/ops/clover_live_mask/jit.h @@ -0,0 +1,34 @@ +#pragma once +#include "base/clover_live_mask.h" +#include "triton/jit/jit.h" + +namespace infini::ops { + +template <> +class Operator + : public CloverLiveMask { + public: + using CloverLiveMask::CloverLiveMask; + + void operator()(const Tensor ids, Tensor previous, int64_t mask_id, + bool refresh, Tensor live) const override { + if (ids.numel() == 0) return; + const triton::jit::Config config{4u, + 1u, + {{"N", ids.numel()}, + {"WIDTH", 1}, + {"MASK", mask_id}, + {"REFRESH", refresh}, + {"STAGE", 0}, + {"BLOCK", 256}}}; + triton::jit::Launch( + "clover_live_rows", ids.device().index(), stream_, nullptr, config, {}, + [&](const triton::jit::Config&) -> std::optional { + return triton::jit::Grid{ + static_cast((ids.numel() + 255) / 256)}; + }, + ids, previous, live, live); + } +}; + +} // namespace infini::ops diff --git a/src/triton/ops/clover_live_rows/jit.py b/src/triton/ops/clover_live_rows/jit.py new file mode 100644 index 000000000..230ba3a84 --- /dev/null +++ b/src/triton/ops/clover_live_rows/jit.py @@ -0,0 +1,39 @@ +import triton +import triton.language as tl + + +@triton.jit +def clover_live_rows_kernel( + A, + B, + C, + D, + N: tl.constexpr, + WIDTH: tl.constexpr, + MASK: tl.constexpr, + REFRESH: tl.constexpr, + STAGE: tl.constexpr, + BLOCK: tl.constexpr, +): + offset = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + valid = offset < N + if STAGE == 0: + current = tl.load(A + offset, valid, 0) + previous = tl.load(B + offset, valid & (not REFRESH), MASK) + live = (current == MASK) | (previous == MASK) | REFRESH + tl.store(C + offset, live, valid) + tl.store(B + offset, current, valid) + elif STAGE == 1: + route = tl.load(A + offset, valid, MASK) + live = tl.load(B + offset // WIDTH, valid, 0) + tl.store(C + offset, tl.where(live != 0, route, MASK), valid) + else: + fresh = tl.load(A + offset, valid, 0) + live = tl.load(B + offset // WIDTH, valid, 0) + cached = tl.load(C + offset, valid & (live == 0), 0) + output = tl.where(live != 0, fresh, cached) + tl.store(D + offset, output, valid) + tl.store(C + offset, output, valid) + + +kernel = clover_live_rows_kernel diff --git a/src/triton/ops/clover_local_memory/jit.h b/src/triton/ops/clover_local_memory/jit.h new file mode 100644 index 000000000..fbbe6011f --- /dev/null +++ b/src/triton/ops/clover_local_memory/jit.h @@ -0,0 +1,62 @@ +#pragma once + +#include + +#include "base/clover_local_memory.h" +#include "triton/jit/jit.h" + +namespace infini::ops::triton::jit { + +template +class CloverLocalMemory : public ::infini::ops::CloverLocalMemory { + public: + using ::infini::ops::CloverLocalMemory::CloverLocalMemory; + + void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor state, const Tensor indices, + const Tensor valid, const Tensor logit_scale, + const Tensor history_bias, const Tensor history_c, + size_t block_size, Tensor out) const override { + if (q.numel() == 0) return; + const Config config{4u, + 1u, + {{"NUM_Q_HEADS", q.size(2)}, + {"NUM_K_HEADS", k.size(2)}, + {"NUM_V_HEADS", v.size(2)}, + {"BLOCK_SIZE", block_size}, + {"NUM_SLOTS", state.size(0)}, + {"KEY_DIM", 128}, + {"VALUE_DIM", 128}, + {"BLOCK_T", 32}, + {"BLOCK_Q", 32}, + {"NUM_QUERY_TILES", 1}, + {"BLOCK_K", 128}, + {"BLOCK_V", 64}}}; + Launch( + "clover_local_memory", q.device().index(), this->stream_, + this->config_ptr_.get(), config, {}, + [&](const Config&) -> std::optional { + return Grid{2, static_cast(indices.numel()), + static_cast(v.size(2))}; + }, + q, k, v, g, state, indices, valid, logit_scale, history_bias, history_c, + out, q.stride(1), q.stride(2), q.stride(3), k.stride(1), k.stride(2), + k.stride(3), v.stride(1), v.stride(2), v.stride(3), g.stride(1), + g.stride(2), state.stride(0), state.stride(1), state.stride(2), + state.stride(3), out.stride(1), out.stride(2), out.stride(3), + 1.0f / std::sqrt(128.0f), 32.0f); + } +}; + +} // namespace infini::ops::triton::jit + +namespace infini::ops { + +template <> +class Operator + : public triton::jit::CloverLocalMemory { + public: + using triton::jit::CloverLocalMemory::CloverLocalMemory; +}; + +} // namespace infini::ops diff --git a/src/triton/ops/clover_local_memory/jit.py b/src/triton/ops/clover_local_memory/jit.py new file mode 100644 index 000000000..8595de10c --- /dev/null +++ b/src/triton/ops/clover_local_memory/jit.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from SGLang PR #46927, commit 058939af24d5cdc44bb4b8852d0ddfab555ed71e. +# Changes: InfiniOps JIT entry point, Hygon launch geometry and pool bounds. +import triton +import triton.language as tl + + +@triton.jit +def clover_local_memory_kernel( + query_ptr, + key_ptr, + value_ptr, + g_ptr, + clean_state_ptr, + cache_indices_ptr, + active_valid_ptr, + softmax_logit_scale_ptr, + softmax_history_bias_ptr, + softmax_history_c_ptr, + output_ptr, + query_stride_t, + query_stride_h, + query_stride_k, + key_stride_t, + key_stride_h, + key_stride_k, + value_stride_t, + value_stride_h, + value_stride_v, + g_stride_t, + g_stride_h, + state_stride_slot, + state_stride_h, + state_stride_v, + state_stride_k, + output_stride_t, + output_stride_h, + output_stride_v, + scale, + softmax_reference_block_size, + NUM_Q_HEADS: tl.constexpr, + NUM_K_HEADS: tl.constexpr, + NUM_V_HEADS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + NUM_SLOTS: tl.constexpr, + KEY_DIM: tl.constexpr, + VALUE_DIM: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_Q: tl.constexpr, + NUM_QUERY_TILES: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_V: tl.constexpr, +): + """Compute the read-only BD-GDN preview output for one fixed-size block. + + Block-local softmax includes the gate-decayed clean readout as one + history bucket. Its logit is history_bias plus + history_c * log(valid_block_size / reference_block_size). + + One program handles one request, one value head, and one value-dimension + tile. Persistent state is read directly in the [slot, H, V, K] + layout and is never written. + """ + value_tile = tl.program_id(0) + query_program = tl.program_id(1) + request_idx = query_program // NUM_QUERY_TILES + query_tile = query_program % NUM_QUERY_TILES + value_head = tl.program_id(2) + + token_offsets = query_tile * BLOCK_Q + tl.arange(0, BLOCK_Q) + local_key_offsets = tl.arange(0, BLOCK_T) + key_offsets = tl.arange(0, BLOCK_K) + value_offsets = value_tile * BLOCK_V + tl.arange(0, BLOCK_V) + token_mask = token_offsets < BLOCK_SIZE + local_key_mask = local_key_offsets < BLOCK_SIZE + key_mask = key_offsets < KEY_DIM + value_mask = value_offsets < VALUE_DIM + token_indices = request_idx * BLOCK_SIZE + token_offsets + + cache_index = tl.load(cache_indices_ptr + request_idx) + request_active = (cache_index >= 0) & (cache_index < NUM_SLOTS) + safe_cache_index = tl.maximum(cache_index, 0) + active = ( + tl.load(active_valid_ptr + token_indices, mask=token_mask, other=0) != 0 + ) & request_active + + # Value heads may share Q/K heads. Clover maps pairs of value heads to + # one of 16 Q/K heads without materializing repeat_interleave. + query_head = value_head // (NUM_V_HEADS // NUM_Q_HEADS) + key_head = value_head // (NUM_V_HEADS // NUM_K_HEADS) + + query_offsets = ( + token_indices[:, None] * query_stride_t + + query_head * query_stride_h + + key_offsets[None, :] * query_stride_k + ) + key_offsets_transposed = ( + (request_idx * BLOCK_SIZE + local_key_offsets)[None, :] * key_stride_t + + key_head * key_stride_h + + key_offsets[:, None] * key_stride_k + ) + value_matrix_offsets = ( + (request_idx * BLOCK_SIZE + local_key_offsets)[:, None] * value_stride_t + + value_head * value_stride_h + + value_offsets[None, :] * value_stride_v + ) + + query = tl.load( + query_ptr + query_offsets, + mask=token_mask[:, None] & key_mask[None, :] & request_active, + other=0.0, + ) + key_transposed = tl.load( + key_ptr + key_offsets_transposed, + mask=key_mask[:, None] & local_key_mask[None, :] & request_active, + other=0.0, + ) + values = tl.load( + value_ptr + value_matrix_offsets, + mask=local_key_mask[:, None] & value_mask[None, :] & request_active, + other=0.0, + ) + + # Match the GDN epsilon and normalize before Tensor Core dot products. + query_f32 = query.to(tl.float32) + query_norm = tl.rsqrt(tl.sum(query_f32 * query_f32, axis=1) + 1.0e-6) + query = (query_f32 * query_norm[:, None]).to(query.dtype) + key_f32 = key_transposed.to(tl.float32) + key_norm = tl.rsqrt(tl.sum(key_f32 * key_f32, axis=0) + 1.0e-6) + key_transposed = (key_f32 * key_norm[None, :]).to(key_transposed.dtype) + + state_offsets = ( + safe_cache_index * state_stride_slot + + value_head * state_stride_h + + value_offsets[None, :] * state_stride_v + + key_offsets[:, None] * state_stride_k + ) + clean_state = tl.load( + clean_state_ptr + state_offsets, + mask=key_mask[:, None] & value_mask[None, :] & request_active, + other=0.0, + ) + + scalar_offsets = token_indices * g_stride_t + value_head * g_stride_h + g = tl.load(g_ptr + scalar_offsets, mask=token_mask & request_active, other=0.0).to( + tl.float32 + ) + local_active = ( + tl.load( + active_valid_ptr + request_idx * BLOCK_SIZE + local_key_offsets, + mask=local_key_mask, + other=0, + ) + != 0 + ) & request_active + + history = tl.dot(query, clean_state.to(query.dtype)) + scores = tl.dot(query, key_transposed) + logit_scale = tl.load(softmax_logit_scale_ptr + value_head).to(tl.float32) + history_bias = tl.load(softmax_history_bias_ptr + value_head).to(tl.float32) + history_c = tl.load(softmax_history_c_ptr + value_head).to(tl.float32) + valid_block_size = tl.maximum(tl.sum(local_active, axis=0), 1.0) + history_bias += history_c * tl.log(valid_block_size / softmax_reference_block_size) + logits = scores * logit_scale + logits = tl.where(local_active[None, :], logits, float("-inf")) + row_max = tl.maximum(tl.max(logits, axis=1), history_bias) + exp_local = tl.exp(logits - row_max[:, None]) + exp_history = tl.exp(history_bias - row_max) + denominator = tl.sum(exp_local, axis=1) + exp_history + local_weights = exp_local / denominator[:, None] + history_weight = exp_history / denominator + local = tl.dot(local_weights.to(values.dtype), values) + output = history * (scale * tl.exp(g) * history_weight)[:, None] + local + output = tl.where(active[:, None], output, 0.0) + + output_offsets = ( + token_indices[:, None] * output_stride_t + + value_head * output_stride_h + + value_offsets[None, :] * output_stride_v + ) + tl.store( + output_ptr + output_offsets, + output, + mask=token_mask[:, None] & value_mask[None, :], + ) + + +kernel = clover_local_memory_kernel diff --git a/src/triton/ops/clover_mask_routes/jit.h b/src/triton/ops/clover_mask_routes/jit.h new file mode 100644 index 000000000..5482d2c3b --- /dev/null +++ b/src/triton/ops/clover_mask_routes/jit.h @@ -0,0 +1,34 @@ +#pragma once +#include "base/clover_mask_routes.h" +#include "triton/jit/jit.h" + +namespace infini::ops { + +template <> +class Operator + : public CloverMaskRoutes { + public: + using CloverMaskRoutes::CloverMaskRoutes; + + void operator()(const Tensor ids, const Tensor live, int64_t num_experts, + Tensor out) const override { + if (ids.numel() == 0) return; + const triton::jit::Config config{4u, + 1u, + {{"N", ids.numel()}, + {"WIDTH", ids.size(1)}, + {"MASK", num_experts}, + {"REFRESH", false}, + {"STAGE", 1}, + {"BLOCK", 256}}}; + triton::jit::Launch( + "clover_live_rows", ids.device().index(), stream_, nullptr, config, {}, + [&](const triton::jit::Config&) -> std::optional { + return triton::jit::Grid{ + static_cast((ids.numel() + 255) / 256)}; + }, + ids, live, out, out); + } +}; + +} // namespace infini::ops diff --git a/src/triton/ops/clover_merge_moe/jit.h b/src/triton/ops/clover_merge_moe/jit.h new file mode 100644 index 000000000..1bef804df --- /dev/null +++ b/src/triton/ops/clover_merge_moe/jit.h @@ -0,0 +1,35 @@ +#pragma once +#include "base/clover_merge_moe.h" +#include "triton/jit/jit.h" + +namespace infini::ops { + +template <> +class Operator + : public CloverMergeMoe { + public: + using CloverMergeMoe::CloverMergeMoe; + + void operator()(const Tensor fresh, const Tensor live, Tensor cache, + Tensor out) const override { + if (fresh.numel() == 0) return; + const triton::jit::Config config{4u, + 1u, + {{"N", fresh.numel()}, + {"WIDTH", fresh.size(1)}, + {"MASK", 0}, + {"REFRESH", false}, + {"STAGE", 2}, + {"BLOCK", 256}}}; + triton::jit::Launch( + "clover_live_rows", fresh.device().index(), stream_, nullptr, config, + {}, + [&](const triton::jit::Config&) -> std::optional { + return triton::jit::Grid{ + static_cast((fresh.numel() + 255) / 256)}; + }, + fresh, live, cache, out); + } +}; + +} // namespace infini::ops diff --git a/src/triton/ops/clover_recurrent_gdn/jit.h b/src/triton/ops/clover_recurrent_gdn/jit.h new file mode 100644 index 000000000..f1e47ac3e --- /dev/null +++ b/src/triton/ops/clover_recurrent_gdn/jit.h @@ -0,0 +1,51 @@ +#pragma once + +#include "base/clover_recurrent_gdn.h" +#include "triton/jit/jit.h" + +namespace infini::ops::triton::jit { + +template +class CloverRecurrentGdn : public ::infini::ops::CloverRecurrentGdn { + public: + using ::infini::ops::CloverRecurrentGdn::CloverRecurrentGdn; + + void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor g, const Tensor beta, const Tensor offsets, + const Tensor initial_indices, const Tensor final_indices, + Tensor state, Tensor out) const override { + const Config config{4u, + 1u, + {{"TOKENS", q.size(1)}, + {"POOL", state.size(0)}, + {"HK", q.size(2)}, + {"HV", v.size(2)}, + {"BV", 16}}}; + Launch( + "clover_recurrent_gdn", q.device().index(), this->stream_, + this->config_ptr_.get(), config, {}, + [&](const Config&) -> std::optional { + return Grid{static_cast(offsets.numel() - 1), + static_cast(v.size(2)), 8}; + }, + q, k, v, g, beta, offsets, initial_indices, final_indices, state, out, + q.stride(1), q.stride(2), k.stride(1), k.stride(2), v.stride(1), + v.stride(2), g.stride(1), g.stride(2), beta.stride(1), beta.stride(2), + state.stride(0), state.stride(1), state.stride(2), state.stride(3), + out.stride(1), out.stride(2)); + } +}; + +} // namespace infini::ops::triton::jit + +namespace infini::ops { + +template <> +class Operator + : public triton::jit::CloverRecurrentGdn { + public: + using triton::jit::CloverRecurrentGdn< + Device::Type::kHygon>::CloverRecurrentGdn; +}; + +} // namespace infini::ops diff --git a/src/triton/ops/clover_recurrent_gdn/jit.py b/src/triton/ops/clover_recurrent_gdn/jit.py new file mode 100644 index 000000000..1019f1d7f --- /dev/null +++ b/src/triton/ops/clover_recurrent_gdn/jit.py @@ -0,0 +1,77 @@ +import triton +import triton.language as tl + + +@triton.jit +def clover_recurrent_gdn_kernel( + Q, + K, + V, + G, + Beta, + Offsets, + Initial, + Final, + State, + Out, + qs, + qh, + ks, + kh, + vs, + vh, + gs, + gh, + bs, + bh, + ss, + sh, + sv, + sk, + os, + oh, + TOKENS: tl.constexpr, + POOL: tl.constexpr, + HK: tl.constexpr, + HV: tl.constexpr, + BV: tl.constexpr, +): + request, head, tile = tl.program_id(0), tl.program_id(1), tl.program_id(2) + begin, end = tl.load(Offsets + request), tl.load(Offsets + request + 1) + initial, final = tl.load(Initial + request), tl.load(Final + request) + if ( + begin >= 0 + and end >= begin + and end <= TOKENS + and initial >= 0 + and initial < POOL + and final >= 0 + and final < POOL + ): + key = tl.arange(0, 128) + value = tile * BV + tl.arange(0, BV) + state = tl.load( + State + initial * ss + head * sh + value[:, None] * sv + key[None, :] * sk + ) + key_head = head // (HV // HK) + for token in range(begin, end): + q = tl.load(Q + token * qs + key_head * qh + key).to(tl.float32) + k = tl.load(K + token * ks + key_head * kh + key).to(tl.float32) + v = tl.load(V + token * vs + head * vh + value).to(tl.float32) + q = q * tl.rsqrt(tl.sum(q * q, 0) + 1.0e-6) * 0.08838834764831845 + k = k * tl.rsqrt(tl.sum(k * k, 0) + 1.0e-6) + decay = tl.exp(tl.load(G + token * gs + head * gh)) + beta = tl.load(Beta + token * bs + head * bh) + state = state * decay + remembered = tl.sum(state * k[None, :], 1) + delta = (v - remembered) * beta + state += delta[:, None] * k[None, :] + y = tl.sum(state * q[None, :], 1) + tl.store(Out + token * os + head * oh + value, y) + tl.store( + State + final * ss + head * sh + value[:, None] * sv + key[None, :] * sk, + state, + ) + + +kernel = clover_recurrent_gdn_kernel diff --git a/src/triton/ops/moe_fused_dense/jit.h b/src/triton/ops/moe_fused_dense/jit.h new file mode 100644 index 000000000..99d8c4123 --- /dev/null +++ b/src/triton/ops/moe_fused_dense/jit.h @@ -0,0 +1,93 @@ +#pragma once + +#include "base/moe_fused_dense.h" +#include "triton/jit/jit.h" + +namespace infini::ops::triton::jit { + +template +class MoeFusedDense : public ::infini::ops::MoeFusedDense { + public: + using ::infini::ops::MoeFusedDense::MoeFusedDense; + + static bool supports(Tensor output, const Tensor hidden, const Tensor w13, + const Tensor w2, const Tensor weights, const Tensor ids, + const Tensor sorted, const Tensor experts, + const Tensor total, size_t block_size, + const std::optional bias1 = std::nullopt, + const std::optional bias2 = std::nullopt) { + return block_size == 16 && !bias1 && !bias2 && + ::infini::ops::MoeFusedDense::supports( + output, hidden, w13, w2, weights, ids, sorted, experts, total, + block_size, bias1, bias2); + } + + size_t workspace_size_in_bytes() const override { + const auto pairs = metadata_[1].size(0) * metadata_[5].size(1); + return pairs * (metadata_[3].size(2) * 2 + metadata_[1].size(1)) * + metadata_[1].element_size(); + } + + void operator()(Tensor output, const Tensor hidden, const Tensor w13, + const Tensor w2, const Tensor weights, const Tensor ids, + const Tensor sorted, const Tensor experts, const Tensor total, + size_t block_size, const std::optional bias1, + const std::optional bias2) const override { + Invocation(output, hidden, w13, w2, weights, ids, sorted, experts, total, + block_size, bias1, bias2); + if (bias1 || bias2) + throw std::invalid_argument( + "Hygon Triton MoE does not support expert biases."); + if (!workspace_ || workspace_size_in_bytes_ < workspace_size_in_bytes()) + throw std::invalid_argument( + "Hygon Triton MoE requires its declared workspace."); + const auto n = hidden.size(0), h = hidden.size(1), i = w2.size(2), + topk = ids.size(1); + const auto pairs = n * topk; + Tensor mid{workspace_, {pairs, 2 * i}, hidden.dtype(), hidden.device()}; + Tensor pair_output{ + static_cast(workspace_) + pairs * 2 * i * hidden.element_size(), + {pairs, h}, + hidden.dtype(), + hidden.device()}; + for (int stage = 0; stage < 3; ++stage) { + const Config config{4u, + 2u, + {{"TOKENS", n}, + {"E", w13.size(0)}, + {"H", h}, + {"I", i}, + {"TOPK", topk}, + {"BM", block_size}, + {"BN", 64}, + {"BK", 64}, + {"STAGE", stage}}}; + Launch( + "moe_fused_dense", hidden.device().index(), this->stream_, nullptr, + config, {}, + [&](const Config&) -> std::optional { + if (stage == 2) + return Grid{static_cast(n), + static_cast((h + 63) / 64)}; + const auto cols = stage == 0 ? 2 * i : h; + return Grid{static_cast(sorted.numel() / block_size), + static_cast((cols + 63) / 64)}; + }, + hidden, w13, w2, weights, ids, sorted, experts, total, mid, + pair_output, output); + } + } +}; + +} // namespace infini::ops::triton::jit + +namespace infini::ops { + +template <> +class Operator + : public triton::jit::MoeFusedDense { + public: + using triton::jit::MoeFusedDense::MoeFusedDense; +}; + +} // namespace infini::ops diff --git a/src/triton/ops/moe_fused_dense/jit.py b/src/triton/ops/moe_fused_dense/jit.py new file mode 100644 index 000000000..e06e63d7e --- /dev/null +++ b/src/triton/ops/moe_fused_dense/jit.py @@ -0,0 +1,103 @@ +import triton +import triton.language as tl + + +@triton.jit +def moe_fused_dense_kernel( + X, + W13, + W2, + Weights, + Ids, + Sorted, + Experts, + Total, + Mid, + Pairs, + Out, + TOKENS: tl.constexpr, + E: tl.constexpr, + H: tl.constexpr, + I: tl.constexpr, # noqa: E741 + TOPK: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, + STAGE: tl.constexpr, +): + if STAGE == 2: + token = tl.program_id(0) + h = tl.program_id(1) * BN + tl.arange(0, BN) + ranks = tl.arange(0, triton.next_power_of_2(TOPK)) + route = tl.load(Ids + token * TOPK + ranks, ranks < TOPK, E) + active = (route >= 0) & (route < E) + values = tl.load( + Pairs + (token * TOPK + ranks[:, None]) * H + h[None, :], + active[:, None] & (h[None, :] < H), + 0, + ).to(tl.float32) + weights = tl.load(Weights + token * TOPK + ranks, ranks < TOPK, 0) + tl.store(Out + token * H + h, tl.sum(values * weights[:, None], 0), h < H) + else: + block = tl.program_id(0) + n = tl.program_id(1) * BN + tl.arange(0, BN) + if block * BM < tl.load(Total): + expert = tl.load(Experts + block) + rows = tl.load(Sorted + block * BM + tl.arange(0, BM)) + valid = rows < TOKENS * TOPK + kk = tl.arange(0, BK) + acc = tl.zeros((BM, BN), tl.float32) + if expert >= 0: + if STAGE == 0: + for start in range(tl.cdiv(H, BK)): + k = start * BK + kk + x = tl.load( + X + (rows[:, None] // TOPK) * H + k[None, :], + valid[:, None] & (k[None, :] < H), + 0, + ) + w = tl.load( + W13 + expert * (2 * I * H) + n[None, :] * H + k[:, None], + (n[None, :] < 2 * I) & (k[:, None] < H), + 0, + ) + acc += tl.dot(x, w) + tl.store( + Mid + rows[:, None] * (2 * I) + n[None, :], + acc, + valid[:, None] & (n[None, :] < 2 * I), + ) + else: + for start in range(tl.cdiv(I, BK)): + k = start * BK + kk + gate = tl.load( + Mid + rows[:, None] * (2 * I) + k[None, :], + valid[:, None] & (k[None, :] < I), + 0, + ).to(tl.float32) + up = tl.load( + Mid + rows[:, None] * (2 * I) + I + k[None, :], + valid[:, None] & (k[None, :] < I), + 0, + ).to(tl.float32) + x = (gate * tl.sigmoid(gate) * up).to(Mid.dtype.element_ty) + w = tl.load( + W2 + expert * (H * I) + n[None, :] * I + k[:, None], + (n[None, :] < H) & (k[:, None] < I), + 0, + ) + acc += tl.dot(x, w) + tl.store( + Pairs + rows[:, None] * H + n[None, :], + acc, + valid[:, None] & (n[None, :] < H), + ) + elif STAGE == 1: + tl.store( + Pairs + rows[:, None] * H + n[None, :], + acc, + valid[:, None] & (n[None, :] < H), + ) + + +kernel = moe_fused_dense_kernel diff --git a/tests/conftest.py b/tests/conftest.py index b54f5ed57..68e083a47 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -47,7 +47,7 @@ def _host_range_profile_compiled(): def pytest_configure(config): - torch.backends.fp32_precision = "tf32" + torch.backends.fp32_precision = "ieee" if torch.version.hip else "tf32" profile_output = config.getoption("--host-range-profile") config._infini_host_range_profile_output = None diff --git a/tests/native/build_clover_workspace.py b/tests/native/build_clover_workspace.py new file mode 100644 index 000000000..56f261634 --- /dev/null +++ b/tests/native/build_clover_workspace.py @@ -0,0 +1,37 @@ +"""Build the native workspace test adapter against an installed InfiniOps SDK.""" + +import os +import pathlib +import subprocess +import sysconfig + +import pybind11 + +root = pathlib.Path(__file__).resolve().parents[2] +prefix = pathlib.Path(os.environ["INFINI_ROOT"]) +cuda = pathlib.Path(os.environ.get("CUDA_HOME", "/opt/dtk/cuda/cuda")) +output = root / "build" / "clover-tests" +output.mkdir(parents=True, exist_ok=True) +subprocess.run( + [ + os.environ.get("CXX", "g++"), + "-std=c++17", + "-O2", + "-shared", + "-fPIC", + str(pathlib.Path(__file__).with_name("clover_workspace_smoke.cpp")), + "-o", + str(output / ("workspace_ops" + sysconfig.get_config_var("EXT_SUFFIX"))), + "-I" + pybind11.get_include(), + "-I" + sysconfig.get_path("include"), + "-I" + str(prefix / "include"), + "-I" + str(root / "src"), + "-I" + str(cuda / "include"), + "-L" + str(prefix / "lib"), + "-linfiniops", + "-linfinirt", + "-Wl,-rpath," + str(prefix / "lib"), + ], + check=True, +) +print(output) diff --git a/tests/native/check_clover_workspace.py b/tests/native/check_clover_workspace.py new file mode 100644 index 000000000..65235b89c --- /dev/null +++ b/tests/native/check_clover_workspace.py @@ -0,0 +1,68 @@ +import infini.ops as ops +import torch +import workspace_ops + +torch.manual_seed(29) +s = torch.cuda.current_stream().cuda_stream +d = "cuda" +dtype = torch.bfloat16 +for n, k in ((1, 8), (32, 8), (64, 8), (32, 3)): + e, h, i, bm = 16, 128, 64, 16 + x = torch.randn(n, h, device=d, dtype=dtype) * 0.2 + w13 = torch.randn(e, 2 * i, h, device=d, dtype=dtype) * 0.1 + w2 = torch.randn(e, h, i, device=d, dtype=dtype) * 0.1 + scores = torch.randn(n, e, device=d).softmax(-1) + weights, ids = scores.topk(k, dim=-1) + weights = weights.contiguous() + ids = ids.int().contiguous() + cap = max(n * k * bm if n * k < e else n * k + e * (bm - 1), bm) + sorted_ids = torch.empty(cap, device=d, dtype=torch.int32) + experts = torch.empty((cap + bm - 1) // bm, device=d, dtype=torch.int32) + total = torch.empty(1, device=d, dtype=torch.int32) + ops.moe_align_block_size(ids, None, e, bm, sorted_ids, experts, total, stream=s) + y = torch.full_like(x, float("nan")) + scratch = torch.empty(n * k * (2 * i + h) * 2, device=d, dtype=torch.uint8) + workspace_ops.moe( + [y, x, w13, w2, weights, ids, sorted_ids, experts, total], bm, scratch, s + ) + ref = torch.zeros_like(x, dtype=torch.float32) + for token in range(n): + for r in range(k): + expert = int(ids[token, r]) + mid = x[token] @ w13[expert].T + act = (torch.nn.functional.silu(mid[:i].float()) * mid[i:].float()).to( + dtype + ) + value = act @ w2[expert].T + ref[token] += value.float() * weights[token, r] + torch.testing.assert_close(y, ref.to(dtype), rtol=0.025, atol=0.003) + print("PASS MoE", n, float((y - ref).abs().max()), flush=True) +q = torch.randn(1, 11, 2, 128, device=d, dtype=dtype) +k = torch.randn_like(q) +v = torch.randn(1, 11, 4, 128, device=d, dtype=dtype) +g = -torch.rand(1, 11, 4, device=d) +beta = torch.rand_like(g) +state = torch.randn(3, 4, 128, 128, device=d) +old = state.clone() +out = torch.empty_like(v) +offsets = torch.tensor([0, 4, 11], device=d, dtype=torch.int32) +init = torch.tensor([0, 1], device=d, dtype=torch.int32) +final = init.clone() +ops.clover_recurrent_gdn(q, k, v, g, beta, offsets, init, final, state, out, stream=s) +qn = q.float() * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) +kn = k.float() * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) +qn = qn.repeat_interleave(2, 2) / 128**0.5 +kn = kn.repeat_interleave(2, 2) +ref = torch.empty_like(v) +for b, (start, end) in enumerate(((0, 4), (4, 11))): + st = old[b].clone() + for t in range(start, end): + st = st * g[0, t].exp()[:, None, None] + remembered = (st * kn[0, t, :, None, :]).sum(-1) + delta = (v[0, t].float() - remembered) * beta[0, t, :, None] + st += delta[:, :, None] * kn[0, t, :, None, :] + ref[0, t] = (st * qn[0, t, :, None, :]).sum(-1) + torch.testing.assert_close(state[b], st, rtol=0.001, atol=0.001) +torch.testing.assert_close(out, ref, rtol=0.02, atol=0.01) +assert torch.equal(state[2], old[2]) +print("PASS recurrent GDN", float((out - ref).abs().max()), flush=True) diff --git a/tests/native/clover_workspace_smoke.cpp b/tests/native/clover_workspace_smoke.cpp new file mode 100644 index 000000000..4b98d5571 --- /dev/null +++ b/tests/native/clover_workspace_smoke.cpp @@ -0,0 +1,41 @@ +#include +#include + +#include "pybind11_utils.h" + +namespace py = pybind11; +using namespace infini::ops; + +PYBIND11_MODULE(workspace_ops, m) { + m.def("moe", [](py::list tensors, size_t block_size, py::object workspace, + uintptr_t stream) { + if (tensors.size() != 9) + throw py::value_error("Expected nine MoE tensors."); + Handle handle; + handle.set_stream(reinterpret_cast(stream)); + auto scratch = TensorFromPybind11Handle(workspace); + handle.set_workspace(scratch.data()); + handle.set_workspace_size_in_bytes(scratch.numel() * + scratch.element_size()); + const auto tensor = [&](int i) { + return TensorFromPybind11Handle(tensors[i]); + }; + MoeFusedDense::Call(handle, + MoeFusedDense::DefaultConfig(Device::Type::kHygon), + tensor(0), tensor(1), tensor(2), tensor(3), tensor(4), + tensor(5), tensor(6), tensor(7), tensor(8), block_size, + std::optional{}, std::optional{}); + }); + m.def("coreset", [](py::object logits, size_t width, double top_p, + py::object out, py::object workspace, uintptr_t stream) { + Handle handle; + handle.set_stream(reinterpret_cast(stream)); + auto w = TensorFromPybind11Handle(workspace); + handle.set_workspace(w.data()); + handle.set_workspace_size_in_bytes(w.numel()); + CloverCoreset::Call(handle, + CloverCoreset::DefaultConfig(Device::Type::kHygon), + TensorFromPybind11Handle(logits), width, top_p, + TensorFromPybind11Handle(out)); + }); +} diff --git a/tests/test_clover_live_rows.py b/tests/test_clover_live_rows.py new file mode 100644 index 000000000..2abe04381 --- /dev/null +++ b/tests/test_clover_live_rows.py @@ -0,0 +1,30 @@ +import infini.ops as ops +import pytest +import torch + + +def test_clover_live_rows(): + if not torch.cuda.is_available() or not torch.version.hip: + pytest.skip("requires Hygon") + if not hasattr(ops, "CloverLiveMask"): + pytest.skip("requires Clover providers") + s = torch.cuda.current_stream().cuda_stream + ids = torch.tensor([151669, 3, 4, 151669], device="cuda", dtype=torch.int64) + previous = torch.empty_like(ids) + live = torch.empty(4, device="cuda", dtype=torch.int8) + ops.clover_live_mask(ids, previous, 151669, True, live, stream=s) + assert live.tolist() == [1, 1, 1, 1] + ids[:] = torch.tensor([5, 3, 4, 151669], device="cuda") + ops.clover_live_mask(ids, previous, 151669, False, live, stream=s) + assert live.tolist() == [1, 0, 0, 1] + routes = torch.arange(16, device="cuda", dtype=torch.int32).view(4, 4) + ops.clover_mask_routes(routes, live, 256, routes, stream=s) + assert routes[1:3].eq(256).all() + fresh = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) + cache = torch.randn_like(fresh) + old = cache.clone() + out = torch.empty_like(fresh) + ops.clover_merge_moe(fresh, live, cache, out, stream=s) + ref = torch.where(live.bool()[:, None], fresh, old) + torch.testing.assert_close(out, ref, rtol=0, atol=0) + assert torch.equal(cache, ref) diff --git a/tests/test_clover_preview.py b/tests/test_clover_preview.py new file mode 100644 index 000000000..c249215d0 --- /dev/null +++ b/tests/test_clover_preview.py @@ -0,0 +1,126 @@ +"""Validate Clover preview outputs and read-only clean-state behavior.""" + +import infini.ops as ops +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _ieee_reference(): + previous = torch.backends.fp32_precision + torch.backends.fp32_precision = "ieee" + try: + yield + finally: + torch.backends.fp32_precision = previous + + +def test_clover_preview_reference_and_state(): + if not torch.cuda.is_available() or not torch.version.hip: + pytest.skip("requires a Hygon Torch environment") + if not hasattr( + ops, "CloverLocalMemory" + ) or not ops.CloverLocalMemory.active_implementation_indices("hygon"): + pytest.skip("requires Hygon Triton Clover providers") + torch.manual_seed(17) + stream = torch.cuda.current_stream().cuda_stream + for width in (1, 8, 32): + for batch in (1, 3): + dtype = torch.bfloat16 + channels = 128 + heads = 4 + x = torch.randn(batch * width, channels, device="cuda", dtype=dtype) + state = torch.randn(batch + 1, channels, 3, device="cuda", dtype=dtype) + state0 = state.clone() + indices = torch.arange(batch, device="cuda", dtype=torch.int32) + weight = torch.randn(channels, 1, 4, device="cuda", dtype=dtype) + right = torch.randn(channels, 1, 3, device="cuda", dtype=dtype) + valid = torch.ones(batch * width, device="cuda", dtype=torch.int8) + out = torch.empty_like(x) + ops.clover_block_conv1d( + x, state, indices, weight, right, valid, width, out, stream=stream + ) + ref = [] + for b in range(batch): + row = x.view(batch, width, channels)[b] + joined = torch.cat((state[b].T, row), 0).float() + y = torch.zeros_like(row, dtype=torch.float32) + for t in range(width): + y[t] = (joined[t : t + 4] * weight[:, 0, :].T.float()).sum(0) + for j in range(1, 4): + if t + j < width: + y[t] += row[t + j].float() * right[:, 0, j - 1].float() + ref.append(torch.nn.functional.silu(y).to(dtype)) + torch.testing.assert_close( + out, torch.stack(ref).view_as(out), rtol=0.02, atol=0.04 + ) + assert torch.equal(state, state0) + q = torch.randn(1, batch * width, 2, 128, device="cuda", dtype=dtype) + k = torch.randn_like(q) + v = torch.randn(1, batch * width, heads, 128, device="cuda", dtype=dtype) + g = -torch.rand(1, batch * width, heads, device="cuda") + s = torch.randn(batch + 1, heads, 128, 128, device="cuda") + s0 = s.clone() + scale = torch.full((heads,), 10.0, device="cuda") + bias = torch.full_like(scale, 5.0) + hc = torch.ones_like(scale) + y = torch.empty_like(v) + ops.clover_local_memory( + q, k, v, g, s, indices, valid, scale, bias, hc, width, y, stream=stream + ) + qn = ( + ( + q.float() + * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) + ) + .to(dtype) + .float() + ) + kn = ( + ( + k.float() + * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) + ) + .to(dtype) + .float() + ) + refs = [] + for b in range(batch): + qi = ( + qn[0, b * width : (b + 1) * width] + .repeat_interleave(2, 1) + .permute(1, 0, 2) + ) + ki = ( + kn[0, b * width : (b + 1) * width] + .repeat_interleave(2, 1) + .permute(1, 2, 0) + ) + vi = v[0, b * width : (b + 1) * width].permute(1, 0, 2).float() + scores = qi @ ki * 10 + history_logit = 5 + torch.log(torch.tensor(width / 32.0)) + weights = torch.softmax( + torch.cat( + ( + scores, + torch.full((heads, width, 1), history_logit, device="cuda"), + ), + -1, + ), + -1, + ) + history = qi @ s[b].to(dtype).float().transpose(-1, -2) + local = weights[:, :, :width].to(dtype).float() @ vi + oi = ( + history + * ( + g[0, b * width : (b + 1) * width].T.exp() + * weights[:, :, -1] + / 128**0.5 + ).unsqueeze(-1) + + local + ) + refs.append(oi.permute(1, 0, 2)) + ref = torch.cat(refs).unsqueeze(0).to(dtype) + torch.testing.assert_close(y, ref, rtol=0.035, atol=0.04) + assert torch.equal(s, s0) diff --git a/tests/test_generate_wrappers.py b/tests/test_generate_wrappers.py index a3efcd026..aa1485fc6 100644 --- a/tests/test_generate_wrappers.py +++ b/tests/test_generate_wrappers.py @@ -1317,3 +1317,15 @@ def test_w4_public_api_selection_is_explicit(tmp_path): module._select_ops_from_config( {}, {name: {"headers": None, "implementations": (0,)}}, "ops.json" ) + + +def test_default_argument_scan_ignores_comment_examples(monkeypatch, tmp_path): + module = _load_generator_module() + header = tmp_path / "sample.h" + header.write_text( + "// bool is_neox=false is GPT-J, true is GPT-NeoX.\n" + "/* float alpha = invalid example, */\n" + "class Sample { public: Sample(bool is_neox, float alpha = 1.0f); };\n" + ) + monkeypatch.setattr(module, "_find_base_header", lambda name: header) + assert module._find_params_with_defaults("sample") == {"alpha": "1.0f"} diff --git a/tests/test_triton_jit_compile.py b/tests/test_triton_jit_compile.py index 4b9b186e0..2ecf00faf 100644 --- a/tests/test_triton_jit_compile.py +++ b/tests/test_triton_jit_compile.py @@ -360,3 +360,27 @@ def test_auto_tuning_cache_rejects_zero_compile_options(monkeypatch, tmp_path, f module.write_auto_tuning_config(identity, config) assert module.read_auto_tuning_config(identity) is None + + +def test_hip_target_keeps_gfx_architecture(monkeypatch): + module, _ = _load_compile_module(monkeypatch) + target = module._create_target( + {"backend": "hip", "architecture": "gfx936", "warp_size": 64} + ) + assert (target.backend, target.arch, target.warp_size) == ("hip", "gfx936", 64) + + +def test_hip_artifacts_without_cuda_scratch_fields(monkeypatch, tmp_path): + module, fake_triton = _load_compile_module(monkeypatch) + fake_triton.compiler.make_backend = lambda target: types.SimpleNamespace( + binary_ext="hsaco" + ) + compiled = types.SimpleNamespace( + metadata=types.SimpleNamespace(name="hip_kernel", num_warps=4, shared=0), + asm={"hsaco": b"hip-binary"}, + ) + prefix = tmp_path / "kernel" + module._write_artifacts(compiled, object(), prefix, "hip-identity") + artifact = module.read_kernel_artifact(prefix, "hip-identity") + assert artifact["binary"] == b"hip-binary" + assert artifact["global_scratch_size"] == artifact["profile_scratch_size"] == 0