From 415ca12afcd1cabd7e19a8b73bea69bbd9a5201d Mon Sep 17 00:00:00 2001 From: project_6 Date: Sun, 16 Aug 2026 16:09:15 +0000 Subject: [PATCH] fix: group_gemm format "TN" + Layer 3 ops_api dispatch from xllm upstream AST chain alignment with upstream_ref/xllm/xllm/core/kernels/ilu/: Layer 5: ixformer::infer (binary .so on device) Layer 4: xllm_kernels/ilu/*.cpp -> calls ixformer::infer (0-diff with upstream) Layer 3: xllm_kernels/ops_api.h+cpp + param.h (NEW from upstream 2719 lines) kernels/kernels.h aggregation header (NEW) Layer 2: xllm_layers/ilu/*.cpp (0-diff with upstream) Layer 1: ix_full_bridge_v2.cpp pybind11 bridge (FIXED) Critical fixes in ix_full_bridge_v2.cpp: - group_gemm format "default" -> "TN" (match upstream ilu/group_gemm.cpp) - fused_moe_forward: pass 3D weights directly instead of .view({-1,...}) - group_gemm output_n: use tokens_per_experts.sum() per upstream convention --- ex_engine/csrc/ix_full_bridge_v2.cpp | 18 +- ex_engine/kernels/kernels.h | 1 + ex_engine/kernels/ops_api.h | 1 + ex_engine/kernels/param.h | 1 + ex_engine/xllm_kernels/kernels.h | 11 + ex_engine/xllm_kernels/ops_api.cpp | 1101 ++++++++++++++++++++ ex_engine/xllm_kernels/ops_api.h | 177 ++++ ex_engine/xllm_kernels/param.h | 1441 ++++++++++++++++++++++++++ 8 files changed, 2747 insertions(+), 4 deletions(-) create mode 120000 ex_engine/kernels/kernels.h create mode 120000 ex_engine/kernels/ops_api.h create mode 120000 ex_engine/kernels/param.h create mode 100644 ex_engine/xllm_kernels/kernels.h create mode 100644 ex_engine/xllm_kernels/ops_api.cpp create mode 100644 ex_engine/xllm_kernels/ops_api.h create mode 100644 ex_engine/xllm_kernels/param.h diff --git a/ex_engine/csrc/ix_full_bridge_v2.cpp b/ex_engine/csrc/ix_full_bridge_v2.cpp index f5396150..576a77be 100644 --- a/ex_engine/csrc/ix_full_bridge_v2.cpp +++ b/ex_engine/csrc/ix_full_bridge_v2.cpp @@ -326,15 +326,21 @@ torch::Tensor ix_moe_expand_input(torch::Tensor input, torch::Tensor ix_group_gemm(torch::Tensor inputs, torch::Tensor weights, torch::Tensor tokens_per_experts, int64_t output_n) { + // Match upstream xllm/core/kernels/ilu/group_gemm.cpp exactly: + // moe_w16a16_group_gemm(output, input, weight, tokens_per_experts, + // dst_to_src=nullopt, bias=nullopt, + // format="TN", persistent=0, + // output_n=tokens_per_experts.sum()) int64_t total_tokens = inputs.size(0); auto output = inputs.new_empty({total_tokens, output_n}); + int64_t gemm_output_n = tokens_per_experts.sum().item(); ixformer::infer::moe_w16a16_group_gemm( output, inputs, weights, tokens_per_experts, /*dst_to_src=*/c10::nullopt, /*bias=*/c10::nullopt, - /*format=*/"default", + /*format=*/"TN", /*persistent=*/0, - output_n); + gemm_output_n); return output; } @@ -382,16 +388,20 @@ torch::Tensor ix_fused_moe_forward( auto expanded = ix_moe_expand_input(hidden_states, src_dst, dst_src, topk); // Step 4: group_gemm (w13: gate_up projection) + // w13 shape: [num_experts, 2*intermediate, hidden] — pass as-is (3D) + // output_n = tokens_per_experts.sum() per upstream convention int64_t intermediate_2x = w13.size(1); - auto gate_up = ix_group_gemm(expanded, w13.view({-1, w13.size(2)}), + int64_t output_n_w13 = expert_sizes_gpu.sum().item(); + auto gate_up = ix_group_gemm(expanded, w13, expert_sizes_gpu, intermediate_2x); // Step 5: silu_and_mul auto activated = ix_silu_and_mul(gate_up); // Step 6: group_gemm (w2: down projection) + // w2 shape: [num_experts, hidden, intermediate] — pass as-is (3D) int64_t hidden_size = w2.size(1); - auto down = ix_group_gemm(activated, w2.view({-1, w2.size(2)}), + auto down = ix_group_gemm(activated, w2, expert_sizes_gpu, hidden_size); // Step 7: moe_combine_result diff --git a/ex_engine/kernels/kernels.h b/ex_engine/kernels/kernels.h new file mode 120000 index 00000000..28c375ea --- /dev/null +++ b/ex_engine/kernels/kernels.h @@ -0,0 +1 @@ +../xllm_kernels/kernels.h \ No newline at end of file diff --git a/ex_engine/kernels/ops_api.h b/ex_engine/kernels/ops_api.h new file mode 120000 index 00000000..c0ff13b6 --- /dev/null +++ b/ex_engine/kernels/ops_api.h @@ -0,0 +1 @@ +../xllm_kernels/ops_api.h \ No newline at end of file diff --git a/ex_engine/kernels/param.h b/ex_engine/kernels/param.h new file mode 120000 index 00000000..527e9687 --- /dev/null +++ b/ex_engine/kernels/param.h @@ -0,0 +1 @@ +../xllm_kernels/param.h \ No newline at end of file diff --git a/ex_engine/xllm_kernels/kernels.h b/ex_engine/xllm_kernels/kernels.h new file mode 100644 index 00000000..30b23bc8 --- /dev/null +++ b/ex_engine/xllm_kernels/kernels.h @@ -0,0 +1,11 @@ +/* Auto-generated aggregation header for xllm::kernel namespace. + * Equivalent to CMake cc_library(NAME kernels HDRS param.h ops_api.h). + * + * AST Layer 3: kernel dispatch interface + * Called by: xllm_layers/ (Layer 2) + * Calls: xllm_kernels/ilu/ (Layer 4) + */ +#pragma once + +#include "param.h" +#include "ops_api.h" diff --git a/ex_engine/xllm_kernels/ops_api.cpp b/ex_engine/xllm_kernels/ops_api.cpp new file mode 100644 index 00000000..40638732 --- /dev/null +++ b/ex_engine/xllm_kernels/ops_api.cpp @@ -0,0 +1,1101 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "ops_api.h" + +#if defined(USE_MLU) +#include "mlu/mlu_ops_api.h" +#elif defined(USE_NPU) +#include "core/kernels/npu/tilelang/tilelang_ops_api.h" +#include "npu/npu_ops_api.h" +#include "triton_npu/torch_api/triton_ops_api.h" +#elif defined(USE_CUDA) +#include "cuda/attention_runner.h" +#include "cuda/cuda_ops_api.h" +#elif defined(USE_ILU) +#include "ilu/ilu_ops_api.h" +#elif defined(USE_MUSA) +#include "cuda/cuda_ops_api.h" +#include "musa/musa_ops_api.h" +#endif + +#include + +#include "common/macros.h" +#include "layers/common/attention_metadata.h" + +namespace xllm::kernel { + +void apply_rotary(RotaryParams& params) { +#if defined(USE_MLU) + mlu::apply_rotary(params.q, + params.k, + params.sin, + params.cos, + params.position_ids, + params.cu_query_lens, + params.interleaved, + params.discrete, + params.dynamic_ntk, + params.max_query_len); +#elif defined(USE_NPU) + npu::apply_rotary( + params.q, params.k, params.cos_sin, params.position_ids.value()); +#elif defined(USE_CUDA) || defined(USE_MUSA) + bool is_neox = !params.interleaved; + torch::Tensor pos_ids; + torch::Tensor cos_sin; + + if (params.position_ids.has_value()) { + // positions is already int64 on CUDA/MUSA (pre-converted in + // ForwardInput::to). + pos_ids = params.position_ids.value().to(torch::kInt64); + } else if (params.cu_query_lens.has_value()) { + auto cu = params.cu_query_lens.value().to(torch::kInt64); + CHECK(cu.numel() >= 2) << "apply_rotary (CUDA): cu_query_lens must have at " + "least 2 elements when " + "position_ids is not provided."; + int64_t seq_len = cu[1].item() - cu[0].item(); + CHECK(seq_len > 0) + << "apply_rotary (CUDA): invalid sequence length inferred from " + "cu_query_lens when position_ids is not provided."; + pos_ids = torch::arange(seq_len, + torch::TensorOptions() + .dtype(torch::kInt64) + .device(params.q.device())) + .contiguous(); + } else { + // When neither position_ids nor cu_query_lens is provided, + // infer sequence length from q tensor and create default position IDs. + // This handles cases like LongCat-Image-Edit where rotary embedding + // is applied uniformly across all sequence positions. + int64_t seq_len = params.q.size(0); + CHECK(seq_len > 0) << "apply_rotary (CUDA): cannot infer valid sequence " + "length from q tensor."; + pos_ids = torch::arange(seq_len, + torch::TensorOptions() + .dtype(torch::kInt64) + .device(params.q.device())) + .contiguous(); + } + + if (params.precomputed_cos_sin.defined()) { + cos_sin = params.precomputed_cos_sin; + } else if (params.cos.defined() && params.sin.defined()) { + const int64_t head_dim = params.cos.size(-1); + const int64_t rot_half = head_dim / 2; + auto cos_sliced = params.cos.contiguous().slice(-1, 0, rot_half); + auto sin_sliced = params.sin.contiguous().slice(-1, 0, rot_half); + cos_sin = torch::cat({cos_sliced, sin_sliced}, -1); + } else if (params.cos_sin.defined()) { + auto cos_sin_vec = params.cos_sin.chunk(4, -1); + auto cos = cos_sin_vec[0]; + auto sin = cos_sin_vec[2]; + cos_sin = torch::cat({cos, sin}, -1); + } else { + LOG(FATAL) << "apply_rotary (CUDA): neither cos_sin nor cos/sin " + "provided; cannot infer cos_sin."; + } + + cuda::rotary_embedding(pos_ids, params.q, params.k, cos_sin, is_neox); +#elif defined(USE_ILU) + torch::Tensor ilu_cos_sin; + if (params.precomputed_cos_sin.defined()) { + ilu_cos_sin = params.precomputed_cos_sin; + } else { + auto cos_sin_vec = params.cos_sin.chunk(4, -1); + ilu_cos_sin = torch::cat({cos_sin_vec[0], cos_sin_vec[2]}, -1); + } + // positions is already int64 on ILU (pre-converted in ForwardInput::to). + torch::Tensor long_position_ids = params.position_ids.value().to(at::kLong); + ilu::apply_rope_pos_ids_cos_sin_cache( + params.q, params.k, ilu_cos_sin, long_position_ids, params.interleaved); +#else + NOT_IMPLEMENTED(); +#endif +} + +void active(ActivationParams& params) { +#if defined(USE_MLU) + mlu::active(params.input, + params.output, + params.bias, + params.cusum_token_count, + params.act_mode, + params.is_gated, + params.start_expert_id, + params.expert_size); +#elif defined(USE_NPU) + params.output = npu::active(params.input, params.act_mode); +#elif defined(USE_CUDA) || defined(USE_MUSA) + cuda::act_and_mul(params.output, params.input, params.act_mode); +#elif defined(USE_ILU) + ilu::act_and_mul(params.output, params.input, params.act_mode); +#else + NOT_IMPLEMENTED(); +#endif +} + +void reshape_paged_cache(ReshapePagedCacheParams& params) { +#if defined(USE_MLU) + mlu::reshape_paged_cache(params.key, + params.value, + params.k_cache, + params.v_cache, + params.slot_mapping, + params.direction); +#elif defined(USE_NPU) + npu::reshape_paged_cache(params.key, + params.value, + params.k_cache, + params.v_cache, + params.slot_mapping); +#elif defined(USE_CUDA) || defined(USE_MUSA) + cuda::reshape_paged_cache(params.slot_mapping, + params.key, + params.value.value_or(torch::Tensor()), + params.k_cache, + params.v_cache.value_or(torch::Tensor())); +#elif defined(USE_ILU) + // auto v_cache = params.v_cache.value_or(torch::Tensor()); + ilu::reshape_paged_cache(params.key, + params.value, + params.k_cache, + params.v_cache, + params.slot_mapping); +#else + NOT_IMPLEMENTED(); +#endif +} + +void reshape_from_cache(ReshapeFromCacheParams& params) { +#if defined(USE_MLU) + mlu::reshape_from_cache(params.key, + params.value, + params.key_cache, + params.value_cache, + params.context_lengths, + params.max_context_len, + params.context_seq_offset, + params.block_tables, + params.cache_seq_offset); +#else + NOT_IMPLEMENTED(); +#endif +} + +void quant_to_paged_cache(ReshapePagedCacheParams& params) { +#if defined(USE_MLU) + CHECK(params.k_cache_scale.has_value()) + << "k_cache_scale is required for quant_to_paged_cache"; + mlu::quant_to_paged_cache(params.key, + params.value, + params.k_cache, + params.v_cache, + params.k_cache_scale.value(), + params.v_cache_scale, + params.slot_mapping); +#else + NOT_IMPLEMENTED(); +#endif +} + +void dequant_from_paged_cache(ReshapeFromCacheParams& params) { +#if defined(USE_MLU) + CHECK(params.key_cache_quant_scale.has_value()) + << "key_cache_quant_scale is required for dequant_from_paged_cache"; + mlu::dequant_from_paged_cache(params.key, + params.value, + params.key_cache, + params.value_cache, + params.key_cache_quant_scale.value(), + params.value_cache_quant_scale, + params.context_lengths, + params.max_context_len, + params.context_seq_offset, + params.block_tables.value(), + params.quant_mode, + params.quant_bit); +#else + NOT_IMPLEMENTED(); +#endif +} + +void fused_layernorm(FusedLayerNormParams& params) { +#if defined(USE_MLU) + mlu::fused_layernorm(params.input, + params.output, + params.residual, + params.weight, + params.beta, + params.bias, + params.quant_scale, + params.residual_out, + params.smooth_quant_scale, + params.normed_out, + params.mode, + params.eps, + params.store_output_before_norm, + params.store_output_after_norm, + params.dynamic_quant); +#elif defined(USE_MUSA) + musa::fused_layernorm(params.input, + params.output, + params.residual, + params.weight, + params.beta, + params.bias, + params.quant_scale, + params.residual_out, + params.smooth_quant_scale, + params.normed_out, + params.mode, + params.eps, + params.store_output_before_norm, + params.store_output_after_norm, + params.dynamic_quant); +#elif defined(USE_NPU) + if (params.residual.has_value()) { + std::tie(params.output, std::ignore, params.residual_out) = + npu::add_rms_norm( + params.input, params.residual.value(), params.weight, params.eps); + } else { + params.output = + npu::rms_norm(params.input, params.weight, params.eps, params.mode); + } +#elif defined(USE_CUDA) || defined(USE_MUSA) + if (params.residual.has_value()) { + cuda::fused_add_rms_norm( + params.input, params.residual.value(), params.weight, params.eps); + params.output = params.input; + params.residual_out = params.residual; + } else { + cuda::rms_norm(params.output, params.input, params.weight, params.eps); + } +#elif defined(USE_ILU) + if (params.residual.has_value()) { + ilu::residual_layer_norm(params.input, + params.output, + params.residual, + params.weight, + params.bias, // residual_bias + params.residual_out, + params.eps); + } else { + ilu::rms_norm(params.output, params.input, params.weight, params.eps); + } +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor matmul(MatmulParams& params) { +#if defined(USE_MLU) + return mlu::matmul( + params.a, params.b, params.bias, params.c, params.alpha, params.beta); +#elif defined(USE_NPU) + return npu::matmul(params.a, params.b, params.bias); +#elif defined(USE_CUDA) || defined(USE_MUSA) + return cuda::matmul(params.a, params.b, params.bias); +#elif defined(USE_ILU) + return ilu::matmul(params.a, params.b, params.bias); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor group_gemm(GroupGemmParams& params) { +#if defined(USE_MLU) + return mlu::group_gemm(params.a, + params.b, + params.token_count, + params.output, + params.a_scale, + params.b_scale, + params.quant_flag, + params.max_dim, + params.trans_a, + params.trans_b, + params.a_quant_bit); +#elif defined(USE_NPU) + std::vector x_list; + std::vector weight_list; + torch::TensorList x_ref; + torch::TensorList weight_ref; + if (params.x_list.has_value()) { + x_ref = params.x_list.value(); + } else { + x_list = {params.a}; + x_ref = x_list; + } + if (params.weight_list.has_value()) { + weight_ref = params.weight_list.value(); + } else { + weight_list = {params.b}; + weight_ref = weight_list; + } + std::optional group_list = params.group_list; + if (!group_list.has_value()) { + group_list = params.token_count; + } + + auto outputs = + npu::apply_npu_grouped_matmul(x_ref, + weight_ref, + params.bias_list, + params.scale_list, + params.offset_list, + params.antiquant_scale_list, + params.antiquant_offset_list, + params.per_token_scale_list, + group_list, + params.activation_input_list, + params.activation_quant_scale_list, + params.activation_quant_offset_list, + params.split_item, + params.group_type, + params.group_list_type, + params.act_type, + params.tuning_config, + params.output_dtype); + return outputs.back(); +#elif defined(USE_ILU) + return ilu::group_gemm(params.a, + params.b, + params.token_count, + params.combine_idx, + params.output); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::tuple moe_active_topk( + MoeFusedTopkParams& params) { +#if defined(USE_MLU) + return mlu::moe_active_topk(params.input, + params.topk, + params.num_expert_group, + params.topk_group, + params.normalize, + params.mask, + params.normed_by, + params.scoring_func, + params.route_scale, + params.e_score_correction_bias); +#elif defined(USE_NPU) + CHECK_EQ(params.scoring_func, "softmax") + << "Only softmax is supported for NPU"; + auto [topk_weights, topk_ids, row_ids] = npu::apply_moe_gating_topk_softmax( + params.input, params.finished, params.topk); + (void)row_ids; + return std::make_tuple(topk_weights, topk_ids); +#elif defined(USE_ILU) + return ilu::moe_active_topk(params.input, + params.topk, + params.num_expert_group, + params.topk_group, + params.normalize, + params.mask, + params.normed_by, + params.scoring_func, + params.route_scale, + params.e_score_correction_bias); +#elif defined(USE_CUDA) || defined(USE_MUSA) + return cuda::moe_fused_topk(params.input, + params.topk, + params.normalize, + params.e_score_correction_bias, + params.scoring_func); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::vector moe_gen_idx(MoeGenIdxParams& params) { +#if defined(USE_MLU) + return mlu::moe_gen_idx(params.expert_id, params.expert_num); +#elif defined(USE_ILU) + return ilu::moe_gen_idx(params.expert_id, params.expert_num); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor moe_expand_input(MoeExpandInputParams& params) { +#if defined(USE_MLU) + return mlu::moe_expand_input(params.input, + params.gather_index, + params.cusum_token_count, + params.start_expert_id, + params.expert_size); +#elif defined(USE_ILU) + return ilu::moe_expand_input( + params.input, params.gather_index, params.combine_idx, params.topk); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor moe_combine_result(MoeCombineResultParams& params) { +#if defined(USE_MLU) + return mlu::moe_combine_result(params.input, + params.reduce_weight, + params.gather_ids, + params.residual, + params.cusum_token_count, + params.start_expert_id, + params.expert_size, + params.bias); +#elif defined(USE_NPU) + std::optional probes = + params.probes.has_value() + ? params.probes + : std::optional(params.reduce_weight); + auto output = npu::apply_npu_moe_token_unpermute(params.input, + params.gather_ids, + probes, + params.padded_mode, + params.restore_shape); + if (params.residual.has_value()) { + output = output + params.residual.value(); + } + return output; +#elif defined(USE_ILU) + return ilu::moe_combine_result(params.input, params.reduce_weight); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor moe_all2all_gen_send_layout( + MoeAll2AllGenSendLayoutParams& params) { +#if defined(USE_MLU) + return mlu::moe_all2all_gen_send_layout(params.token_count, params.nrank); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::vector moe_all2all_gen_gather_index( + MoeAll2AllGenGatherIndexParams& params) { +#if defined(USE_MLU) + return mlu::moe_all2all_gen_gather_index( + params.token_num, params.pad_num, params.return_cusum_token_count); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::vector moe_all2all_create(MoeAll2AllCreateParams& params) { +#if defined(USE_MLU) + return mlu::moe_all2all_create(params.dispatch_token_byte, + params.combine_token_byte, + params.max_expert_num, + params.max_token_num, + params.rank, + params.nrank, + params.device); +#else + NOT_IMPLEMENTED(); +#endif +} + +void moe_all2all_init(MoeAll2AllInitParams& params) { +#if defined(USE_MLU) + mlu::moe_all2all_init(params.handle, params.all_exchange_info, params.device); +#else + NOT_IMPLEMENTED(); +#endif +} + +void moe_all2all_dispatch(MoeAll2AllDispatchParams& params) { +#if defined(USE_MLU) + mlu::moe_all2all_dispatch(params.handle, + params.token_byte, + params.token_num, + params.send_layout, + params.send_token_num, + params.recv_layout, + params.recv_token_num, + params.send_token, + params.recv_token); +#else + NOT_IMPLEMENTED(); +#endif +} + +void moe_all2all_combine(MoeAll2AllCombineParams& params) { +#if defined(USE_MLU) + mlu::moe_all2all_combine(params.handle, + params.token_byte, + params.token_num, + params.send_src_layout, + params.send_dst_layout, + params.send_token, + params.recv_token); +#else + NOT_IMPLEMENTED(); +#endif +} + +void moe_all2all_destroy(MoeAll2AllDestroyParams& params) { +#if defined(USE_MLU) + mlu::moe_all2all_destroy(params.handle, params.device); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::tuple scaled_quantize( + ScaledQuantizeParams& params) { +#if defined(USE_MLU) + return mlu::scaled_quantize(params.x, + params.smooth, + params.zero, + params.token_count, + params.gather_index, + params.gather_index_start_position, + params.output, + params.output_scale, + params.act_mode, + params.active_coef, + params.is_gated, + params.quant_type); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor scaled_matmul(ScaledMatmulParams& params) { +#if defined(USE_MLU) + return mlu::scaled_matmul(params.a, + params.b, + params.a_scale, + params.b_scale, + params.output_dtype, + params.bias, + params.c, + params.act_mode, + params.quant_bit_size, + params.alpha, + params.beta, + params.use_hp_active, + params.a_quant_bit_size, + params.a_calib, + params.b_calib, + params.output); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor apply_top_k_top_p(TopKPParams& params) { +#if defined(USE_MLU) + return mlu::apply_top_k_top_p( + params.logits, params.temperatures, params.top_k, params.top_p); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor random_sample(RandomSampleParams& params) { +#if defined(USE_MLU) + return mlu::random_sample(params.logits); +#elif defined(USE_CUDA) + return cuda::random_sample(params.logits); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor rejection_sample(RejectionSampleParams& params) { +#if defined(USE_MLU) + return mlu::rejection_sample(params.draft_token_ids, + params.num_draft_tokens, + params.cu_num_draft_tokens, + params.draft_probs, + params.target_probs, + params.bonus_token_ids, + params.uniform_rand, + params.uniform_probs, + params.max_spec_len); +#else + NOT_IMPLEMENTED(); +#endif +} + +void masked_indexer_select_paged_kv(MaskedIndexerSelectPagedKVParams& params) { +#if defined(USE_MLU) + mlu::masked_indexer_select_paged_kv(params.query, + params.k_cache, + params.weights, + params.kv_cache_block_table, + params.cu_seq_q_lens, + params.cu_seq_k_lens, + params.k_context_lens, + params.k_cache_block_table, + params.is_prefill, + params.index_topk, + params.kv_cache_block_size, + params.softmax_scale, + params.q_scale, + params.k_scale_cache, + params.sparse_block_table, + params.sparse_context_lens); +#else + NOT_IMPLEMENTED(); +#endif +} + +void gather_split(GatherSplitParams& params) { +#if defined(USE_MLU) + mlu::gather_split(params.input, + params.gather_index, + params.valid_token_num, + params.output_head, + params.output_tail); +#else + NOT_IMPLEMENTED(); +#endif +} + +void fused_mla_q(FusedMlaQParams& params) { +#if defined(USE_MLU) + mlu::fused_mla_q(params.q, + params.output, + params.output_scale, + params.output_norm, + params.gamma, + params.smooth_quant_scale, + params.weight_b, + params.weight_b_scale, + params.weight_c, + params.sin, + params.cos, + params.position_id, + params.quant_mode, + params.eps, + params.interleaved); +#else + NOT_IMPLEMENTED(); +#endif +} + +void fused_mla_kv(FusedMlaKVParams& params) { +#if defined(USE_MLU) + mlu::fused_mla_kv(params.input_kv, + params.sin, + params.cos, + params.position_id, + params.gamma, + params.kv_cache, + params.kv_cache_scale, + params.slot_mapping, + params.cache_bs_id, + params.cache_seq_offset, + params.quant_mode, + params.is_paged_cache, + params.eps, + params.interleaved); +#else + NOT_IMPLEMENTED(); +#endif +} + +void fused_indexer_q(FusedIndexerQParams& params) { +#if defined(USE_MLU) + mlu::fused_indexer_q(params.input_q, + params.output, + params.output_scale, + params.w_q, + params.w_q_scale, + params.hadamard_matrix, + params.sin, + params.cos, + params.position_id, + params.quant_mode, + params.interleaved, + params.rope_at_front); +#else + NOT_IMPLEMENTED(); +#endif +} + +void fused_indexer_k(FusedIndexerKParams& params) { +#if defined(USE_MLU) + mlu::fused_indexer_k(params.x, + params.wk, + params.wproj, + params.sin_table, + params.cos_table, + params.position_id, + params.slot_mapping, + params.head_weights, + params.k_cache, + params.k_cache_scale, + params.hadamard_matrix, + params.interleaved, + params.gamma, + params.beta, + params.eps); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor l2_norm(torch::Tensor& x, double eps) { +#if defined(USE_NPU) + return npu::npu_l2norm_last_dim(x, eps); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::tuple +moe_init_routing_v2(MoeInitRoutingV2Params& params) { +#if defined(USE_NPU) + return npu::apply_npu_moe_init_routing_v2(params.x, + params.expert_idx, + params.scale, + params.offset, + params.active_num, + params.expert_capacity, + params.expert_num, + params.drop_pad_mode, + params.expert_tokens_num_type, + params.expert_tokens_num_flag, + params.quant_mode, + params.active_expert_range, + params.row_idx_type); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::tuple fp8_scaled_quantize( + Fp8ScaledQuantizeParams& params) { +#if defined(USE_CUDA) + return cuda::fp8_scaled_quantize(params.input, params.output, params.scale); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::pair fused_gdn_gating( + FusedGdnGatingParams& params) { +#if defined(USE_NPU) + return npu::tilelang::fused_gdn_gating(params.A_log, + params.a, + params.b, + params.dt_bias, + params.beta, + params.threshold); + // return npu::npu_fused_gdn_gating(params.A_log, + // params.a, + // params.b, + // params.dt_bias, + // params.beta, + // params.threshold); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::pair fused_recurrent_gated_delta_rule( + FusedRecurrentGatedDeltaRuleParams& params) { +#if defined(USE_NPU) + return npu::npu_fused_recurrent_gated_delta_rule( + params.q, + params.k, + params.v, + params.g, + params.beta, + params.scale, + params.initial_state, + params.inplace_final_state, + params.cu_seqlens, + params.ssm_state_indices, + params.num_accepted_tokens, + params.use_qk_l2norm_in_kernel); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor fp8_scaled_matmul(Fp8ScaledMatmulParams& params) { +#if defined(USE_CUDA) + auto out_2d = cuda::fp8_scaled_matmul(params.a, + params.b, + params.a_scale, + params.b_scale, + params.output_dtype, + params.bias, + params.output); + + // Auto reshape output if original input shape is provided + if (params.input_shape.has_value()) { + auto out_shape = params.input_shape.value(); + out_shape.back() = params.b.size(0); + return out_2d.view(out_shape); + } + return out_2d; +#else + LOG(FATAL) << "fp8_scaled_matmul is only supported on CUDA"; + return torch::Tensor(); +#endif +} + +void static_scaled_fp8_quant(StaticScaledFp8QuantParams& params) { +#if defined(USE_CUDA) + cuda::static_scaled_fp8_quant(params.output, params.input, params.scale); +#else + LOG(FATAL) << "static_scaled_fp8_quant is only supported on CUDA"; +#endif +} + +// Fused RMSNorm + Static FP8 Quantization +torch::Tensor rms_norm_static_fp8_quant(RmsNormStaticFp8QuantParams& params) { +#if defined(USE_CUDA) + auto org_shape = params.input.sizes().vec(); + auto hidden_size = params.input.size(-1); + + // Flatten input to 2D. Use reshape to support non-contiguous tensors. + auto input_2d = params.input.reshape({-1, hidden_size}); + + torch::Tensor output = + torch::empty({input_2d.size(0), hidden_size}, + input_2d.options().dtype(torch::kFloat8_e4m3fn)); + + // Call fused kernel + cuda::rms_norm_static_fp8_quant( + output, input_2d, params.weight, params.scale, params.epsilon); + + return output.reshape(org_shape); +#else + LOG(FATAL) << "rms_norm_static_fp8_quant is only supported on CUDA"; + return torch::Tensor(); +#endif +} + +std::tuple fused_add_rms_norm_static_fp8_quant( + FusedAddRmsNormStaticFp8QuantParams& params) { +#if defined(USE_CUDA) + auto org_shape = params.input.sizes().vec(); + auto hidden_size = params.input.size(-1); + + // Flatten tensors to 2D. Use reshape to support non-contiguous tensors. + auto input_2d = params.input.reshape({-1, hidden_size}); + auto residual_2d = params.residual.reshape({-1, hidden_size}); + + torch::Tensor output = + torch::empty({input_2d.size(0), hidden_size}, + input_2d.options().dtype(torch::kFloat8_e4m3fn)); + + // Call fused kernel (residual is updated in-place) + cuda::fused_add_rms_norm_static_fp8_quant(output, + input_2d, + residual_2d, + params.weight, + params.scale, + params.epsilon); + + // Reshape outputs + auto output_reshaped = output.reshape(org_shape); + auto residual_reshaped = residual_2d.reshape(org_shape); + + return std::make_tuple(output_reshaped, residual_reshaped); +#else + LOG(FATAL) << "fused_add_rms_norm_static_fp8_quant is only supported on CUDA"; + return std::make_tuple(torch::Tensor(), torch::Tensor()); +#endif +} + +torch::Tensor causal_conv1d_update(CausalConv1dUpdateParams& params) { +#if defined(USE_NPU) + if (params.conv_state_indices.has_value()) { + CHECK(params.conv_state_indices.value().is_contiguous()) + << "causal_conv1d_update: conv_state_indices must be contiguous."; + } + return npu::npu_causal_conv1d_update_v2(params.x, + params.conv_state, + params.weight, + params.activation, + params.bias, + params.conv_state_indices, + params.query_start_loc, + params.max_query_len, + params.pad_slot_id, + params.block_idx_last_scheduled_token, + params.initial_state_idx, + params.validate_data); + +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor gated_layer_norm(GatedLayerNormParams& params) { +#if defined(USE_NPU) + return npu::layer_norm_fwd(params.x, + params.weight, + params.bias, + params.eps, + params.z, + params.group_size, + params.norm_before_gate, + params.is_rms_norm); +#elif defined(USE_MLU) + return mlu::gated_layer_norm(params.x, + params.weight, + params.bias, + params.eps, + params.z, + params.group_size, + params.norm_before_gate); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::pair partial_rotary_embedding( + PartialRotaryEmbeddingParams& params) { +#if defined(USE_NPU) + return npu::apply_npu_partial_rotary_embedding(params.positions, + params.query, + params.key, + params.head_size, + params.rotary_dim, + params.cos_sin_cache, + params.is_neox_style); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::tuple +fused_qkvzba_split_reshape_cat(FusedQkvzbaSplitReshapeParams& params) { +#if defined(USE_NPU) + return npu::npu_fused_qkvzba_split_reshape_cat(params.mixed_qkvz, + params.mixed_ba, + params.num_heads_qk, + params.num_heads_v, + params.head_qk, + params.head_v); +#else + NOT_IMPLEMENTED(); +#endif +} + +void gemma_rms_norm(GemmaRMSNormParams& params) { +#if defined(USE_NPU) + npu::npu_gemma_rms_norm( + params.x, params.gamma, params.epsilon, params.rstd_out, params.norm_out); +#elif defined(USE_MLU) + mlu::gemma_rms_norm(params.x, params.gamma, params.epsilon, params.norm_out); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::tuple +split_qkv_rmsnorm_mrope(SplitQkvRmsnormMropeParams& params) { +#if defined(USE_NPU) + return npu::tilelang::split_qkv_rmsnorm_mrope(params.qkvg, + params.q_weight, + params.k_weight, + params.cos_sin, + params.gather_pattern, + params.eps, + params.num_q_heads, + params.num_kv_heads, + params.head_size); +#else + NOT_IMPLEMENTED(); +#endif +} + +bool has_split_qkv_rmsnorm_mrope_specialization(int64_t num_q_heads, + int64_t num_kv_heads, + int64_t head_size) { +#if defined(USE_NPU) + return npu::tilelang::has_split_qkv_rmsnorm_mrope_specialization( + num_q_heads, num_kv_heads, head_size); +#else + return false; +#endif +} + +torch::Tensor build_split_qkv_rmsnorm_mrope_gather_pattern( + int64_t rope_dim, + const std::vector& mrope_section, + bool is_interleaved, + const torch::Device& device) { +#if defined(USE_NPU) + return npu::tilelang::build_split_qkv_rmsnorm_mrope_gather_pattern( + rope_dim, mrope_section, is_interleaved, device); +#else + NOT_IMPLEMENTED(); +#endif +} + +std::pair chunk_gated_delta_rule( + ChunkGatedDeltaRuleParams& params) { +#if defined(USE_NPU) + return npu::npu_chunk_gated_delta_rule(params.q, + params.k, + params.v, + params.g, + params.beta, + params.scale, + params.initial_state, + params.output_final_state, + params.cu_seqlens, + params.head_first, + params.use_qk_l2norm_in_kernel); +#else + NOT_IMPLEMENTED(); +#endif +} + +torch::Tensor recurrent_gated_delta_rule( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + torch::Tensor& state, + const std::optional& beta, + const std::optional scale, + const std::optional& actual_seq_lengths, + const std::optional& ssm_state_indices, + const std::optional& num_accepted_tokens, + const std::optional& g, + const std::optional& gk) { +#if defined(USE_NPU) + return npu::npu_recurrent_gated_delta_rule(query, + key, + value, + state, + beta, + scale, + actual_seq_lengths, + ssm_state_indices, + num_accepted_tokens, + g, + gk); +#else + NOT_IMPLEMENTED(); +#endif +} +} // namespace xllm::kernel diff --git a/ex_engine/xllm_kernels/ops_api.h b/ex_engine/xllm_kernels/ops_api.h new file mode 100644 index 00000000..f355eef7 --- /dev/null +++ b/ex_engine/xllm_kernels/ops_api.h @@ -0,0 +1,177 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include "param.h" + +namespace xllm::kernel { + +static const std::string kActModeSilu = "silu"; +static const std::string kActModeGelu = "gelu"; +static const std::string kActModeQuickGelu = "quick_gelu"; +static const std::string kActModeSwish = "swish"; + +void apply_rotary(RotaryParams& params); + +void active(ActivationParams& params); + +void reshape_paged_cache(ReshapePagedCacheParams& params); + +void reshape_from_cache(ReshapeFromCacheParams& params); + +// Quantize and store KV cache to paged cache (INT8 quantization) +// Only supported on MLU backend +void quant_to_paged_cache(ReshapePagedCacheParams& params); + +// Dequantize KV cache from paged cache (INT8 to FP16/BF16) +// Only supported on MLU backend +void dequant_from_paged_cache(ReshapeFromCacheParams& params); + +void fused_layernorm(FusedLayerNormParams& params); + +torch::Tensor matmul(MatmulParams& params); + +torch::Tensor group_gemm(GroupGemmParams& params); + +std::tuple moe_active_topk( + MoeFusedTopkParams& params); + +std::vector moe_gen_idx(MoeGenIdxParams& params); + +torch::Tensor moe_expand_input(MoeExpandInputParams& params); + +torch::Tensor moe_combine_result(MoeCombineResultParams& params); + +torch::Tensor moe_all2all_gen_send_layout( + MoeAll2AllGenSendLayoutParams& params); + +std::vector moe_all2all_gen_gather_index( + MoeAll2AllGenGatherIndexParams& params); + +std::vector moe_all2all_create(MoeAll2AllCreateParams& params); + +void moe_all2all_init(MoeAll2AllInitParams& params); + +void moe_all2all_dispatch(MoeAll2AllDispatchParams& params); + +void moe_all2all_combine(MoeAll2AllCombineParams& params); + +void moe_all2all_destroy(MoeAll2AllDestroyParams& params); + +std::tuple scaled_quantize( + ScaledQuantizeParams& params); + +torch::Tensor scaled_matmul(ScaledMatmulParams& params); + +torch::Tensor apply_top_k_top_p(TopKPParams& params); + +torch::Tensor random_sample(RandomSampleParams& params); + +torch::Tensor rejection_sample(RejectionSampleParams& params); + +void masked_indexer_select_paged_kv(MaskedIndexerSelectPagedKVParams& params); + +void gather_split(GatherSplitParams& params); + +void fused_mla_q(FusedMlaQParams& params); + +void fused_mla_kv(FusedMlaKVParams& params); + +void fused_indexer_q(FusedIndexerQParams& params); + +void fused_indexer_k(FusedIndexerKParams& params); + +// L2 normalization along the last dimension +torch::Tensor l2_norm(torch::Tensor& x, double eps = 1e-6); + +// TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx + moe_expand_input +// (and token_count/cusum outputs) on other backends. +std::tuple +moe_init_routing_v2(MoeInitRoutingV2Params& params); + +// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format +// Returns: (quantized_output, scale) +std::tuple fp8_scaled_quantize( + Fp8ScaledQuantizeParams& params); + +// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels +// Performs: c = (a @ b.T) with scales applied +torch::Tensor fp8_scaled_matmul(Fp8ScaledMatmulParams& params); + +// Static scaled FP8 quantization helper +// Quantizes input tensor to FP8 using a pre-computed scale factor +void static_scaled_fp8_quant(StaticScaledFp8QuantParams& params); + +// Fused RMSNorm + Static FP8 Quantization +// These fused operations combine RMSNorm and FP8 quantization to reduce memory +// bandwidth by avoiding the intermediate write-back to global memory. + +// Fused RMSNorm + Static FP8 Quantization +// Returns: FP8 quantized output tensor +torch::Tensor rms_norm_static_fp8_quant(RmsNormStaticFp8QuantParams& params); + +// Fused Add + RMSNorm + Static FP8 Quantization (with residual) +// Returns: tuple of (FP8 quantized output, updated residual) +std::tuple fused_add_rms_norm_static_fp8_quant( + FusedAddRmsNormStaticFp8QuantParams& params); + +std::pair fused_gdn_gating( + FusedGdnGatingParams& params); + +std::pair fused_recurrent_gated_delta_rule( + FusedRecurrentGatedDeltaRuleParams& params); + +torch::Tensor causal_conv1d_update(CausalConv1dUpdateParams& params); + +torch::Tensor gated_layer_norm(GatedLayerNormParams& params); + +std::pair partial_rotary_embedding( + PartialRotaryEmbeddingParams& params); + +std::tuple +fused_qkvzba_split_reshape_cat(FusedQkvzbaSplitReshapeParams& params); + +void gemma_rms_norm(GemmaRMSNormParams& params); + +std::tuple +split_qkv_rmsnorm_mrope(SplitQkvRmsnormMropeParams& params); + +bool has_split_qkv_rmsnorm_mrope_specialization(int64_t num_q_heads, + int64_t num_kv_heads, + int64_t head_size); + +torch::Tensor build_split_qkv_rmsnorm_mrope_gather_pattern( + int64_t rope_dim, + const std::vector& mrope_section, + bool is_interleaved, + const torch::Device& device); + +std::pair chunk_gated_delta_rule( + ChunkGatedDeltaRuleParams& params); + +torch::Tensor recurrent_gated_delta_rule( + const torch::Tensor& query, + const torch::Tensor& key, + const torch::Tensor& value, + torch::Tensor& state, + const std::optional& beta, + const std::optional scale, + const std::optional& actual_seq_lengths, + const std::optional& ssm_state_indices, + const std::optional& num_accepted_tokens, + const std::optional& g, + const std::optional& gk); +} // namespace xllm::kernel diff --git a/ex_engine/xllm_kernels/param.h b/ex_engine/xllm_kernels/param.h new file mode 100644 index 00000000..9c96c837 --- /dev/null +++ b/ex_engine/xllm_kernels/param.h @@ -0,0 +1,1441 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include + +namespace xllm::layer { +struct AttentionMetadata; +} // namespace xllm::layer + +namespace xllm::kernel { + +// Note: add default values for optional parameters in the struct definition + +// Rotary embedding parameters +struct RotaryParams { + // Query tensor. First dimension is total_seq_len (T). + // Will be reshaped to [T, -1] and concatenated with k before applying rotary + // embedding. Head size must be between 2 and 256. + torch::Tensor q; + // Key tensor. First dimension must match q.size(0) (total_seq_len). + // Will be reshaped to [T, -1] and concatenated with q before applying rotary + // embedding. + torch::Tensor k; + // Sin cache tensor for rotary embedding. Shape: + // - [rope_seqlen, rope_dim] if dynamic_ntk=false + // - [batch_size, rope_seqlen, rope_dim] if dynamic_ntk=true + // rope_dim must be between 2 and head_size, and must be even. + // rope_dim is extracted as sin.size(-1) and used to reshape qk tensor. + torch::Tensor sin; + // Cos cache tensor for rotary embedding. Same shape as sin. + // The rope_seqlen-stride must equal to sin's rope_seqlen-stride. + torch::Tensor cos; + // Precomputed cos_sin tensor. Not used in current MLU implementation + // (rope.cpp). + torch::Tensor cos_sin; + // Pre-formatted cos_sin cache for kernels that need [cos_half, sin_half] + // layout (CUDA, MUSA, ILU). Avoids chunk/cat operations per layer. + torch::Tensor precomputed_cos_sin; + // Optional position IDs tensor. Type must be int32. + // Shape: [total_seqlen] if discrete=true, or [batch_size] if discrete=false. + // If discrete=true, position_ids must be provided. + std::optional position_ids; + // Cumulative query lengths tensor. Type must be int32, must be contiguous. + // Required in pack mode (when q/k are 3D). Size should be [batch_size + 1]. + // Note: In current MLU implementation, this is always passed to underlying + // API. + std::optional cu_query_lens; + // Whether to use interleaved rotary embedding pattern. + bool interleaved; + // Whether to use discrete position mode. If true, position_ids must be + // provided and have shape [total_seqlen]. If false, position_ids can be None + // or have shape [batch_size]. + bool discrete; + // Whether to use dynamic NTK (Neural Tangent Kernel) scaling. + // If true, sin and cos caches must have batch dimension. + // Note: Current MLU implementation hardcodes this to false when calling + // underlying API, so dynamic_ntk=true may not be fully supported. + bool dynamic_ntk = false; + // Maximum query length. In pad mode (4D input), must equal to input.size(1). + // Must be less than or equal to rope_seqlen if not using discrete + // position_ids. + int64_t max_query_len; +}; + +// Activation parameters +struct ActivationParams { + // Input tensor. Must be contiguous, dimension >= 2. + // Last dimension is in_channel, which must be > 0. + // If is_gated=true, in_channel must be even. + torch::Tensor input; + // Output tensor. Must be contiguous, dimension >= 2. + // Must have same attributes (device, dtype) as input. + // Only supports stride in dim(-2), stride(-1) must be 1. + // Shape: [total_tokens, inner_size] where inner_size = in_channel/2 if + // is_gated else in_channel. + torch::Tensor output; + // Optional bias tensor, only used for MoE activation. + // If provided, cusum_token_count must also be provided. + // Shape: [expert_size, in_channel]. Must be contiguous. + std::optional bias; + // Optional cumulative token count tensor. Type should be int32. + // Required when bias is provided. Must be contiguous. + // Size: [num_expert + 1], where num_expert = size(0) - 1. + std::optional cusum_token_count; + // Activation mode string. Must be one of: "silu", "gelu", "quick_gelu", + // "swish". + // - "silu": SiLU activation (Swish-1) + // - "gelu": GELU activation + // - "quick_gelu": Quick GELU with coefficient 1.702 + // - "swish": Swish activation + std::string act_mode; + // Whether to use gated activation. If true, input's last dimension + // (in_channel) must be even, and output's inner_size will be in_channel/2. + bool is_gated; + // Starting expert ID for MoE activation. Used when processing multiple + // experts. + int64_t start_expert_id = 0; + // Expert size for MoE activation. Used when bias is provided. + // Bias tensor shape must be [expert_size, in_channel]. + int64_t expert_size = 0; +}; + +// Reshape paged cache parameters +struct ReshapePagedCacheParams { + // Key tensor from context. Shape: [num_tokens, num_heads, head_dim]. + // Last two dimensions must be contiguous: stride(-1)==1, + // stride(-2)==head_dim. Must have same device and dtype as k_cache and + // v_cache. + torch::Tensor key; + // Optional value tensor from context. Shape: [num_tokens, num_heads, + // head_dim]. If provided, v_cache must also be provided (and vice versa). + // Last two dimensions must be contiguous: stride(-1)==1, + // stride(-2)==head_dim. Must have same device and dtype as other tensors. + std::optional value; + // Key cache tensor in paged format. Shape: [num_blocks, num_heads, + // block_size, head_dim]. Must be contiguous. Must have same device and dtype + // as key and value. + torch::Tensor k_cache; + // Optional value cache tensor in paged format. Shape: [num_blocks, num_heads, + // block_size, head_dim]. If provided, value must also be provided (and vice + // versa). Must be contiguous. Must have same device and dtype as other + // tensors. + std::optional v_cache; + // Slot mapping tensor. Shape: [num_tokens]. Type must be int32. + // Maps each token to its corresponding slot in the cache. Must be contiguous. + // Must have same device as key. + torch::Tensor slot_mapping; + // Direction flag: false = CONTEXT2CACHE (copy from context to cache), + // true = CACHE2CONTEXT (copy from cache to context). + bool direction = false; + // Optional scale tensor for quantized key cache. Shape: [num_blocks, + // num_heads, block_size]. Dtype: float32. Required when using INT8 + // quantization. + std::optional k_cache_scale; + // Optional scale tensor for quantized value cache. Shape: [num_blocks, + // num_heads, block_size]. Dtype: float32. Required when using INT8 + // quantization. + std::optional v_cache_scale; +}; + +// ReshapeFromCacheParams describes parameters for gathering and flattening +// KV (Key/Value) cached data from a possibly paged or non-contiguous storage +// format into a contiguous tensor. +struct ReshapeFromCacheParams { + // Target tensor to store reshaped key values. Shape: [total_length, head_num, + // head_size]. Dtype: float32, float16, bfloat16, int8. + torch::Tensor key; + // Optional target tensor to store reshaped value values. If provided, + // value_cache must also be provided. Shape: [total_length, head_num, + // head_size]. Dtype: float32, float16, bfloat16, int8. + std::optional value; + // Source tensor containing cached key values. + // Shape: + // - Linear mode: [max_batch_size, head_num, cache_mem_len, head_size] + // - Paged mode: [total_blocks, head_num, block_size, head_size] + // Dtype: float32, float16, bfloat16, int8. + torch::Tensor key_cache; + // Optional source tensor containing cached value values. If provided, value + // must also be provided. Shape: + // - Linear mode: [max_batch_size, head_num, cache_mem_len, head_size] + // - Paged mode: [total_blocks, head_num, block_size, head_size] + // Dtype: float32, float16, bfloat16, int8. + std::optional value_cache; + // 1D tensor representing the lengths of each batch context. + // Shape: [batch_size]. Dtype: int32. + torch::Tensor context_lengths; + // Maximum context length that can be processed at once. + // Used for memory allocation and bounds checking. + int64_t max_context_len; + // Optional 1D tensor with per-context sequence offsets. + // If provided, applies a shift offset for each context's beginning location. + // Shape: [batch_size]. Dtype: int32. Default: None. + std::optional context_seq_offset; + // Optional tensor containing the block indices for each batch. + // Shape: + // - Linear mode: [batch_size, 1] + // - Paged mode: [batch_size, max_blocks] + // Dtype: int32. Default: None (linear mode). + std::optional block_tables; + // Optional 1D tensor representing the cache sequence offset for each batch. + // Used for slicing key and value cache starts in memory. + // Shape: [batch_size]. Dtype: int32. Default: None. + std::optional cache_seq_offset; + + // ========== Quantization parameters (for dequant_from_paged_cache) + // ========== Optional scale tensor for quantized key cache. Shape: + // [num_blocks, num_heads, block_size] or [num_heads, head_dim]. Dtype: + // float32. Required when dequantizing INT8 cache. + std::optional key_cache_quant_scale; + // Optional scale tensor for quantized value cache. + // Shape: [num_blocks, num_heads, block_size] or [num_heads, head_dim]. + // Dtype: float32. Required when dequantizing INT8 cache. + std::optional value_cache_quant_scale; + // Quantization mode: 0 for per-channel, 1 for per-token. Default: 1. + int64_t quant_mode = 1; + // Quantization bit size. Default: 8 (INT8). + int64_t quant_bit = 8; +}; + +// Fused layer norm parameters +struct FusedLayerNormParams { + // Input tensor. Dimension must be >= 2. Last dimension is hidden_size. + // Last dimension must be contiguous: stride(-1) == 1. + // Must have same device and dtype as residual, weight, beta, bias, + // residual_out, normed_out. + torch::Tensor input; + // Output tensor. Must have same shape as input. + // If inplace (input.data_ptr() == output.data_ptr()), strides must also be + // the same. Must have same device as input, smooth_quant_scale, quant_scale. + torch::Tensor output; + // Optional residual tensor. Must have same shape as input. + // If provided, must have same device and dtype as input. + std::optional residual; + // Weight tensor (gamma). Shape: [hidden_size]. Must be contiguous. + // Required for both layernorm and rmsnorm modes. + // Must have same device and dtype as input. + torch::Tensor weight; + // Optional beta tensor. Shape: [hidden_size]. Must be contiguous. + // Required for layernorm mode, not used in rmsnorm mode. + // If provided, must have same dtype as weight. + std::optional beta; + // Optional bias tensor. Shape: [hidden_size]. Must be contiguous. + // Must have same device and dtype as input. + std::optional bias; + // Optional quantization scale tensor. Type must be float. + // Shape: [hidden_size] (1D) or [head, headdim] (2D). + // - 1D: per-channel quantization, input will be flattened to 2D + // - 2D: only supported for rmsnorm mode, input must be dim >= 3, + // shape must be [head, headdim], residual and bias not supported + // If dynamic_quant=true, this must be provided. + std::optional quant_scale; + // Optional residual output tensor. Used when store_output_before_norm=true. + // Not supported when both bias and residual are not provided. + // Must have same device and dtype as input. + std::optional residual_out; + // Optional smooth quantization scale tensor. Type must be float. + // Used when dynamic_quant=true. Will be flattened to 1D. + // Must have same device as input. + std::optional smooth_quant_scale; + // Optional normalized output tensor. Used when store_output_after_norm=true. + // Only supported when dynamic_quant=true. + // Must have same device and dtype as input. + std::optional normed_out; + // Normalization mode. Must be "layernorm" or "rmsnorm". + // - "layernorm": requires both weight (gamma) and beta + // - "rmsnorm": only requires weight (gamma), beta is not used + std::string mode; + // Epsilon value for numerical stability in normalization computation. + double eps; + // Whether to store output before normalization to residual_out. + // Not supported when both bias and residual are not provided. + bool store_output_before_norm = false; + // Whether to store output after normalization to normed_out. + // Only supported when dynamic_quant=true. + bool store_output_after_norm = false; + // Whether to use dynamic quantization. If true, quant_scale must be provided. + // When true, uses per-token quantization scheme; otherwise uses per-channel + // if quant_scale provided. + bool dynamic_quant = false; +}; + +// Matmul parameters +struct MatmulParams { + // Left input tensor A. Must be 2D or 3D. Must have same dimension as b. + // Must have same dtype as b. + // For 2D: shape [M, K], output will be [M, N] where N = b.size(-1) + // For 3D: shape [batch, M, K], output will be [batch, M, N] + // If input dtype is int8 or fp8, c must be provided to determine output + // dtype. + torch::Tensor a; + // Right input tensor B. Must be 2D or 3D. Must have same dimension as a. + // Must have same dtype as a. + // For 2D: shape [K, N], output will be [M, N] where M = a.size(-2) + // For 3D: shape [batch, K, N], output will be [batch, M, N] + torch::Tensor b; + // Optional bias tensor. Will be added to the matrix multiplication result. + std::optional bias; + // Optional output tensor C. Can be used to specify output dtype and + // accumulate result. If input dtype is int8 or fp8, c or dtype must be + // provided to determine output dtype. If provided, result will be: output = + // alpha * (a @ b) + beta * c + std::optional c; + // Scaling factor for matrix multiplication result. Default: 1.0 + // Result: alpha * (a @ b) + beta * c (if c provided) + double alpha = 1.0; + // Scaling factor for tensor c (if provided). Default: 0.0 + // Result: alpha * (a @ b) + beta * c (if c provided) + double beta = 0.0; +}; + +struct GroupGemmParams { + // Input activation tensor. + // Shape: 2D [M, K] if trans_a==false; [K, M] if trans_a==true. + // Must be contiguous. Dtype: float16, bfloat16, or float32. + // Must have same dtype and device as b, output. + torch::Tensor a; + // Weight tensor. + // If trans_b is true, shape is (num_experts, N, K) or (N, K); + // if trans_b is false, shape is (num_experts, K, N) or (K, N). + // Must be contiguous. Dtype and device must match a, output. + torch::Tensor b; + // Per-expert token count tensor. + // Shape: 1D [num_experts]. Type must be int32. + // Controls number of tokens processed per group/expert. + torch::Tensor token_count; + // Output tensor. + // Shape: [num_experts, N] or [num_experts, N, K]. num_experts = + // token_count.size(0). Must be contiguous. Dtype and device must match a. + torch::Tensor output; + // Optional scale tensor for a (input activation), used in quantized mode. + // Shape depends on quantization granularity. + std::optional a_scale; + // Optional scale tensor for b (weight), used in quantized mode. + // Shape depends on quantization granularity. + std::optional b_scale; + // Optional quantization config flag list. + // Used to control per-expert weight quantization mode. + std::optional> quant_flag; + // Maximum workspace dimension (e.g., maximum tokens per expert allowed). + // Used for configuring inner kernel workspace. + int64_t max_dim; + // Whether to transpose a: + // false: [M, K] (default); true: [K, M]. + bool trans_a; + // Whether to transpose b: + // false: [K, N] (default); true: [N, K]. + bool trans_b; + // Quantization bit-width for input a. + // Set -1 to disable quantization. + int64_t a_quant_bit; + // ========== Torch NPU related parameters ========== + // Optional input tensor list for grouped matmul. + // If provided, this overrides `a` for NPU backend. + // Each tensor shape: [M, K] (or [K, M] if trans_a is true). + std::optional x_list; + // Optional weight tensor list for grouped matmul. + // If provided, this overrides `b` for NPU backend. + // Each tensor shape: [K, N] or [N, K] depending on trans_b. + std::optional weight_list; + // Optional bias list. Used in quantized or fused-activation paths. + std::optional bias_list; + // Optional scale list for quantized weights. + std::optional scale_list; + // Optional offset list for quantized weights. + std::optional offset_list; + // Optional anti-quantization scale list. + std::optional antiquant_scale_list; + // Optional anti-quantization offset list. + std::optional antiquant_offset_list; + // Optional per-token scale list. + std::optional per_token_scale_list; + // Optional group list for NPU grouped matmul. + // If group_list_type == 0: values are cumsum of group sizes. + // If group_list_type == 1: values are per-group sizes. + std::optional group_list; + // Optional activation input list for fused activation. + std::optional activation_input_list; + // Optional activation quantization scale list. + std::optional activation_quant_scale_list; + // Optional activation quantization offset list. + std::optional activation_quant_offset_list; + // Optional split item for grouped matmul. + // Common value is 2 for gated MLP (gate + up). + std::optional split_item = 2; + // Optional group type for grouped matmul. + // 0 indicates grouping along the M axis (row-wise). + std::optional group_type = 0; + // Optional group list type for grouped matmul. + // 0: cumsum of group sizes; 1: per-group sizes. + std::optional group_list_type = 1; + // Optional activation type for fused activation. + std::optional act_type; + // Optional tuning configuration for NPU kernel. + c10::OptionalIntArrayRef tuning_config; + // Optional output dtype for NPU kernel. + std::optional output_dtype; + // ========== Torch ILU related parameters ========== + // Inverse mapping of gather_idx. + // Shape: [expand_token_num]. + // Dtype: int32. + std::optional combine_idx; +}; + +struct MoeFusedTopkParams { + // Input tensor. + // Shape: [*, num_mask, num_expert] (e.g., [batch, num_mask, num_expert]). + // Dtype: float32, float16, bfloat16. + // Must be contiguous. + torch::Tensor input; + // Optional finished mask for NPU gating topk softmax. + // Shape should be broadcastable to input's leading dims. + // If not provided, all tokens are considered active. + std::optional finished; + // Number of top-k experts to select per token. + // Constraint: 0 < topk <= num_expert. + int64_t topk; + // Number of expert groups for group-limited top-k selection. + // If > 1, mask must be None, and num_expert % num_expert_group == 0. + int64_t num_expert_group; + // Maximum selected experts per group. + // Constraint: 0 < topk_group <= num_expert_group. + int64_t topk_group; + // Whether to renormalize expert weights after top-k selection. + bool normalize; + // Optional mask tensor. + // Shape: [1, ..., 1, num_mask, num_expert] (leading dims must be 1). + // Dtype must match input. + // Must be contiguous. + std::optional mask; + // Normalization logic after top-k selection. + // For softmax: "topk_logit" or "softmax_logit". + // For sigmoid: "topk_logit" or "sigmoid_logit". + std::string normed_by; + // Scoring function for expert selection. + // Supported: "softmax", "sigmoid". + std::string scoring_func; + // Route scaling factor applied to routing scores. + double route_scale; + // Optional expert score correction bias. + // Shape: [num_expert]. + // Dtype: float32, float16, or bfloat16. + // Must be contiguous. + std::optional e_score_correction_bias; +}; + +struct MoeGenIdxParams { + // The input tensor stores the expert id of each token. + // Shape: [num_tokens, topk]. + // Dtype: int32. + torch::Tensor expert_id; + // Expert number. + // Must be >= 0. + int64_t expert_num; +}; + +struct MoeExpandInputParams { + // Input tensor to be expanded. + // Shape: [token_num, hidden_size]. + // Dtype: int8, float, half, or bfloat16. + torch::Tensor input; + // Index tensor for gather operation. + // Shape: [expand_token_num]. + // Dtype: int32. + torch::Tensor gather_index; + // Optional prefix sum of token count per expert. + // Shape: [num_experts + 1]. + // Dtype: int32. + // If provided, adjusts gather range for each expert. + std::optional cusum_token_count; + // Starting expert id to process. + // Must be >= 0. + int64_t start_expert_id; + // Number of experts to process in this call. + // Must be >= 0. + int64_t expert_size; + // ========== Torch ILU related parameters ========== + // Inverse mapping of gather_idx. + // Shape: [expand_token_num]. + // Dtype: int32. + torch::Tensor combine_idx; + // topk for moe + int topk; +}; + +struct MoeCombineResultParams { + // Expert output tensor to be combined. + // Shape: [num_tokens * topk, hidden_size]. + // - Must be contiguous. + // - Dtype: float32, float16, or bfloat16. + // - This is the concatenated output from all experts, not yet reordered back + // to the original sequence order. + torch::Tensor input; + // Router/gating weights tensor. Used for weighted combination of expert + // outputs. Shape: [num_tokens, topk]. + // - Must be contiguous at last dimension. + // - Dtype: float32. + // - Constraint: reduce_weight.numel() == input.size(0). + torch::Tensor reduce_weight; + // Gather index tensor that maps combined output to original token positions. + // Shape: [num_tokens * topk]. + // - Must be contiguous. + // - Dtype: int32. + // - Corresponds to permutation/scatter indices for reordering expert outputs. + torch::Tensor gather_ids; + // Optional probes tensor for NPU token unpermute. + // If provided, used as probe weights in unpermute kernel. + // Shape: [num_tokens, topk]. + std::optional probes; + // Whether the permuted tokens are padded (NPU token unpermute). + bool padded_mode = false; + // Optional restore shape for NPU token unpermute. + c10::OptionalIntArrayRef restore_shape = c10::nullopt; + // Optional residual connection input. + // Shape: [num_tokens, hidden_size]. + // - Must have same shape and dtype as output if provided. + // - Must be contiguous if provided. + // - Default: std::nullopt (no residual). + std::optional residual; + // Optional cumulative token count for expert assignment. + // Shape: [num_experts + 1] or deduced by expert_size. + // - Must be contiguous if provided. + // - Dtype: int32. + // - Used to infer num_expert or assist calculation in some kernels. + std::optional cusum_token_count; + // Starting expert ID + // - Must be >= 0. + // - Used to mark the offset of current experts being processed (for + // sharding). + int64_t start_expert_id = 0; + // Number of experts processed in this step. + // - If cusum_token_count not given, num_expert is set to this value. + // - If cusum_token_count given, deduced num_expert must satisfy: + // num_expert >= start_expert_id + expert_size + int64_t expert_size = 0; + // Optional bias tensor. + // WARNING: Bias addition is NOT supported in current implementation. + // Always keep as std::nullopt unless bias support is added in the future. + std::optional bias; +}; + +struct MoeAll2AllGenSendLayoutParams { + // Expert token count tensor. + // Shape: [expert_num]. + // Dtype: int32. + // Each element represents the number of tokens assigned to each expert. + torch::Tensor token_count; + // Number of ranks (processes) participating in All2All. + // Must be >= 0. + int64_t nrank; +}; + +struct MoeAll2AllGenGatherIndexParams { + // The table that indicates the relationship of token for each Expert Parallel + // part. Shape: [rank_num, expert_num], where rank_num is the number of + // devices in Expert Parallel, and expert_num is the number of experts handled + // by each device. Dtype: int32. + torch::Tensor token_num; + // The max token count for each rank (used for padding). + // Dtype: int32. Must be >= 0. + int64_t pad_num; + // Whether to return the cusum_token_count tensor. + // If true, cusum_token_count will be returned. + bool return_cusum_token_count = false; +}; + +struct MoeAll2AllCreateParams { + // Byte size of a single token for dispatch All-to-All operation. + // Each token to be dispatched requires this many bytes. + int64_t dispatch_token_byte; + // Byte size of a single token for combine All-to-All operation. + // Each token to be combined requires this many bytes. + int64_t combine_token_byte; + // Maximum number of experts participating in the All-to-All operation. + // (Sets the upper bound for how many experts can be involved. + int64_t max_expert_num; + // Maximum number of tokens to be processed. + // Upper bound on the total batch size in tokens for the operation. + int64_t max_token_num; + // Rank ID of the current process in the distributed group, within [0, + // nrank-1]. Identifies this process within the world group. + int64_t rank; + // Total number of processes in the distributed group. + // Used for collective communication context and split assignment. + int64_t nrank; + // The current compute device to be used、 + // default to CPU + torch::Device device = torch::Device(torch::kCPU); +}; + +struct MoeAll2AllInitParams { + // communication backend handle for All-to-All operation. + // obtained from moe_all2all_create. + int64_t handle; + // CPU tensor containing aggregated exchange information from all nrank + // processes. + torch::Tensor all_exchange_info; + // The current compute device to be used + // default to CPU + torch::Device device = torch::Device(torch::kCPU); +}; + +struct MoeAll2AllDispatchParams { + // Communication backend handle for All-to-All operation. + // Obtained from moe_all2all_create. + int64_t handle; + // Byte size of a single token. + int64_t token_byte; + // Number of tokens to be processed in the current operation. + int64_t token_num; + // Offset and token count for each rank. + // The token_count is generated by moe_gen_idx. + // Shape: [nrank, 2]. Type: int32. + torch::Tensor send_layout; + // Number of tokens to send to each expert. + // Shape: [max_expert_num]. Type: int32. + torch::Tensor send_token_num; + // Offset and token count from peer ranks. + // Shape: [nrank, 2]. Type: int32. + torch::Tensor recv_layout; + // Expected number of tokens to receive from each expert. + // Shape: [max_expert_num]. Type: int32. + torch::Tensor recv_token_num; + // Optional tensor containing tokens to dispatch. + // If not provided, defaults to dispatch_send created by moe_all2all_create. + std::optional send_token; + // Optional buffer for receiving tokens. + // If not provided, defaults to dispatch_recv created by moe_all2all_create. + std::optional recv_token; +}; + +struct MoeAll2AllCombineParams { + // communication backend handle for All-to-All operation. + // obtained from moe_all2all_create. + int64_t handle; + // Byte size of a single token. + int64_t token_byte; + // The number of tokens to receive. + int64_t token_num; + // The offset and token count for each rank, output from + // Shape: [nrank, 2], + // Type: int32. + torch::Tensor send_src_layout; + // The expected receive pattern from peer ranks. + // Shape: [nrank, 2], + // Type: int32. + torch::Tensor send_dst_layout; + // Optional tensor containing the tokens to dispatch. If not provided, + // defaults to combine_send created by moe_all2all_create. + std::optional send_token; + // Optional buffer for receiving tokens. If not provided, + // defaults to combine_recv created by moe_all2all_create. + std::optional recv_token; +}; + +struct MoeAll2AllDestroyParams { + // communication backend handle for All-to-All operation. + // obtained from moe_all2all_create. + int64_t handle; + // The current compute device to be used + // default to CPU + torch::Device device = torch::Device(torch::kCPU); +}; + +// Per token smooth quantize parameters +// Note: Current MLU implementation uses "dynamic_per_token" quantization mode. +struct ScaledQuantizeParams { + // Input tensor to quantize. Dimension must be >= 2. + // Must be continuous between 0 and -2 dimensions (can be flattened to 2D). + // If gather_index or token_count has value, x must be 2D. + // Must have same device as other tensors. + torch::Tensor x; + // Smooth quantization scale tensor (corresponds to x_scale in underlying + // API). Shape constraints depend on quantization mode and other parameters. + // - If token_count has value: shape [token_count.size(0), + // x.size(-1)/(1+is_gated)] + // - If is_gated: smooth.size(-1) * 2 == x.size(-1) + // - Otherwise: smooth.size(-1) == x.size(-1) + // Must be contiguous if provided. Must have same device as x. + torch::Tensor smooth; + // Zero point tensor. Must be None (not supported in current implementation). + std::optional zero; + // Optional token count tensor when quantizing MoE group gemm inputs. + // If provided, x must be 2D and smooth.size(0) must equal + // token_count.size(0). Must be contiguous if provided. Must have same device + // as x. + std::optional token_count; + // Optional gather index tensor when quantizing MoE group gemm inputs. Shape: + // [output_tokens]. If provided, x must be 2D. Output shape will be adjusted: + // output_shape[0] = gather_index.size(0). If gather_index_start_position is + // provided, gather_index must also be provided. Must be contiguous if + // provided. Must have same device as x. + std::optional gather_index; + // Optional gather index start position tensor when quantizing MoE group gemm + // inputs. Only used if gather_index is provided. Must be contiguous if + // provided. Must have same device as x. + std::optional gather_index_start_position; + // Optional output tensor when quantizing MoE group gemm inputs. + // Type must be int8 (kChar), float8_e4m3fn, or float8_e5m2. + // Dimension must be >= 2. Must be continuous between 0 and -2 dimensions. + // Shape constraints: + // - If !gather_index && !is_gated: output.sizes() == x.sizes() + // - If is_gated: output.size(-1) * 2 == x.size(-1) + // - If gather_index: output_shape[0] = gather_index.size(0) + // If not provided, will be allocated automatically with quant_type. + // Must have same device as x. + std::optional output; + // Optional output scale tensor. + // Used in dynamic_per_token quantization mode. + // Shape: x.sizes()[0:-1] (same as x except last dimension removed). + // If gather_index provided: shape[0] = gather_index.size(0). + // Must be flattenable to 1D with numel == output_flat.size(0). + // If not provided, will be allocated automatically with float32 dtype. + // Must have same device as x. + std::optional output_scale; + // Activation mode. Must be one of: "none", "gelu", "silu", "swish". + // Default: "none". If "none", is_gated will be set to false automatically. + // If "silu", active_coef will be set to 1.0 automatically. + std::string act_mode = "none"; + // Activation coefficient. Default: 1.0. + // If act_mode == "silu", this will be set to 1.0 automatically. + double active_coef = 1.0; + // Whether to use gated activation. Default: false. + // If act_mode == "none", this will be set to false automatically. + // If true, output's last dimension will be x.size(-1) / 2. + bool is_gated = false; + // Quantization output data type. Default: torch::kChar (int8). + // Supported: torch::kChar (int8), torch::kFloat8_e4m3fn, torch::kFloat8_e5m2. + torch::ScalarType quant_type = torch::kChar; +}; + +// Scaled matmul parameters +// Note: Current MLU implementation only supports: +// - smooth_quant algorithm +// - w8a8 quantization (quant_bit_size=8, a_quant_bit_size=8) +// - trans_a=false, trans_b=true (hardcoded) +struct ScaledMatmulParams { + // Input tensor A. Shape: [M, K]. Must be contiguous. + // Output shape will be [M, N] where N = b.size(0). + // Must have same device as other tensors. + torch::Tensor a; + // Weight tensor B. Shape: [K, N]. Will be transposed (trans_b=true). + // Must be contiguous. Must have same device as other tensors. + torch::Tensor b; + // Optional scale tensor for A. Shape: 1D or 2D. Must be contiguous or have + // stride (1, m). + // - 1D: per-token quantization layout + // - 2D: group-wise quantization layout + // Note: In current MLU implementation (scaled_matmul.cpp), a_scale is + // required. + std::optional a_scale; + // Scale tensor for B. Shape: 1D or 2D. Must be contiguous or have stride (1, + // n). Determines quantization layout: + // - 1D: per-channel quantization + // - 2D: per-block (if b_scale.size(0) < b.size(0)) or group-wise quantization + // Must be contiguous. Must have same device as other tensors. + torch::Tensor b_scale; + // Output data type. Must be torch::kFloat16 (half) or torch::kBFloat16. + torch::ScalarType output_dtype; + // Optional bias tensor. Will be added to the matrix multiplication result. + // Must be contiguous. Must have same device as other tensors. + std::optional bias; + // Optional tensor C for accumulation. Result: alpha * (a @ b) + beta * c. + // Must be contiguous. Must have same device as other tensors. + std::optional c; + // Activation mode. Default: "none". Supported: "none", "silu", "gelu". + // If "silu", act_coef will be set to 1.0 automatically. + std::string act_mode = "none"; + // Quantization bit size for B (weight). Default: 8. + // Current implementation only supports 8 (w8a8 quantization). + // Supported values: 4, 8. + int64_t quant_bit_size = 8; + // Scaling factor for matrix multiplication result. Default: 1.0 + // Result: alpha * (a @ b) + beta * c (if c provided) + double alpha = 1.0; + // Scaling factor for tensor c (if provided). Default: 1.0 + // Result: alpha * (a @ b) + beta * c (if c provided) + double beta = 1.0; + // Whether to use high precision activation computation. Default: false + // If true, uses high precision; otherwise uses fast computation. + bool use_hp_active = false; + // Quantization bit size for A (activation). Default: -1. + // Current implementation only supports 8 (w8a8 quantization). + // Supported values: -1 (no quantization), 4, 8. + int64_t a_quant_bit_size = -1; + // Optional calibration tensor for A. Used for flat_quant and svd_quant + // algorithms. Must be contiguous. Must have same device as other tensors. + std::optional a_calib; + // Optional calibration tensor for B. Used for flat_quant and svd_quant + // algorithms. Must be contiguous. Must have same device as other tensors. + std::optional b_calib; + // Optional output tensor. Shape: [M, N] where M = a.size(0), N = b.size(0). + // If not provided, will be allocated automatically with output_dtype. + // Must have same device as other tensors. + std::optional output; +}; + +// Top-K and Top-P sampling parameters +struct TopKPParams { + // Input logits tensor. Shape: [batch_size, vocab_size]. Type must be float32. + // Must be contiguous. Will be converted to float32 if needed. + // If both top_k and top_p are not defined, logits will be returned directly. + torch::Tensor logits; + // Temperature tensor for scaling logits. Shape: [batch_size]. + // Must be contiguous. Will be moved to same device as logits. + torch::Tensor temperatures; + // Optional top-k values tensor. Type will be converted to int32. + // Must be contiguous. Will be moved to same device as logits. + torch::Tensor top_k; + // Optional top-p (nucleus sampling) values tensor. + // Must be contiguous. Will be moved to same device as logits. + torch::Tensor top_p; +}; + +// Random sample parameters +struct RandomSampleParams { + // Input tensor of probabilities for sampling. + // Must be 2-dimensional: [batch_size, vocab_size] + torch::Tensor logits; +}; + +// Rejection sampling parameters for speculative decoding +struct RejectionSampleParams { + // Candidate draft token indices to be verified. + // Shape: [total_draft_tokens]. Dtype: int32. + // total_draft_tokens equals cu_num_draft_tokens[batch_size - 1]. + torch::Tensor draft_token_ids; + // Number of draft tokens for each sequence in the batch. + // Shape: [batch_size]. Dtype: int32. + torch::Tensor num_draft_tokens; + // Accumulated number of draft tokens in each batch. + // Shape: [batch_size]. Dtype: int32. + torch::Tensor cu_num_draft_tokens; + // Probability distributions of the draft model. + // Shape: [total_draft_tokens, vocab_size]. + // Dtype: float32, float16, or bfloat16. + std::optional draft_probs; + // Probability distributions of the target model. + // Shape: [total_draft_tokens, vocab_size]. + // Dtype: float32, float16, or bfloat16. + torch::Tensor target_probs; + // Bonus token indices to be selected when all draft tokens are accepted. + // Shape: [batch_size]. Dtype: int32. + torch::Tensor bonus_token_ids; + // Random probabilities for acceptance threshold comparison. + // Shape: [total_draft_tokens]. Dtype: float32. + // Used to compare with selected_target_probs / selected_draft_probs. + torch::Tensor uniform_rand; + // Random probabilities for resampling (recovery) calculation. + // Shape: [total_draft_tokens, vocab_size]. Dtype: float32. + torch::Tensor uniform_probs; + // The maximum number of draft tokens in the batch (max value in + // num_draft_tokens). + int32_t max_spec_len; +}; + +// Masked indexer select paged KV cache parameters +struct MaskedIndexerSelectPagedKVParams { + // Query tensor. Must have same dtype as k_cache (bfloat16, half, or int8). + // - Prefill mode: 3D [total_seq_q, head_num, head_size], head_num must be 64 + // - Decode mode: 4D [batch_num, len_q, head_num, head_size], head_num must be + // 64 Does not need to be contiguous + torch::Tensor query; + // Key cache tensor in paged format. Shape: [num_blocks, 1, block_size, + // head_dim]. Dim(1) must be 1. Must be contiguous. Must have same dtype as + // query. + torch::Tensor k_cache; + // Attention weights tensor. Dtype must be bfloat16 or float32. Must be + // contiguous. + torch::Tensor weights; + // Key cache block table. Shape: [batch_num, k_cache_max_blkn]. Type: int32. + // Must be contiguous. + std::optional k_cache_block_table; + // Cumulative sequence lengths for queries. Type: int32. Must be contiguous. + // Required in prefill mode, not used in decode mode. + std::optional cu_seq_q_lens; + // Cumulative sequence lengths for keys. + std::optional cu_seq_k_lens; + // Key context lengths tensor. Shape: [batch_num]. Type: int32. Must be + // contiguous. + std::optional k_context_lens; + // KV cache block table. Shape: [batch_num, kv_cache_max_blkn]. Type: int32. + // Must be contiguous. + torch::Tensor kv_cache_block_table; + // Whether this is prefill phase (true) or decode phase (false). + // Affects query shape and whether cu_seq_q_lens is used. + bool is_prefill; + // Number of top-k indices to select. Must be >= 0. + int64_t index_topk; + // KV cache block size. + int64_t kv_cache_block_size; + // Softmax scaling factor for attention computation. + double softmax_scale; + // Query quantization scale tensor. Must be contiguous. + // - Required (numel > 0) when query dtype is int8 or fp8 + // - Must be empty (numel == 0) when query dtype is bfloat16 or half + std::optional q_scale; + // Key cache quantization scale tensor. Must be contiguous. + // - Required (numel > 0) when k_cache dtype is int8 or fp8 + // - Must be empty (numel == 0) when k_cache dtype is bfloat16 or half + std::optional k_scale_cache; + // New sparse block table output tensor. Must be contiguous. + // - Prefill mode: 2D [total_seq_q, kv_cache_max_blkn] + // - Decode mode: 3D [batch_num, seq_q, kv_cache_max_blkn] + torch::Tensor sparse_block_table; + // New sparse block table output tensor. Shape: [batch_num] (prefill) or + // [batch_num] (decode). Type: int32. Must be contiguous. + torch::Tensor sparse_context_lens; +}; + +struct GatherSplitParams { + // Input tensor. Shape: (token_num, input_size). + // Dtype: int8, float32, float16, or bfloat16. + torch::Tensor input; + // Gather index tensor. Shape: (token_num). + // Dtype: int32. + // Used to select valid tokens from the input tensor. + torch::Tensor gather_index; + // Number of valid tokens tensor. Shape: (1). + // Dtype: int32. + // Its first element is the actual valid token count: valid_token_num = + // valid_token_num[0].item(). + torch::Tensor valid_token_num; + // Output tensor for the "head" split. Shape: (token_num, size_0). + // Dtype: same as input. + // Holds the gathered and split tokens for the first size_0 elements of each + // token. + torch::Tensor output_head; + // Optional output tensor for the "tail" split. Shape: (token_num, input_size + // - size_0). Dtype: same as input. If provided, holds the gathered and split + // tokens for the remaining elements after size_0. + // Pass empty tensor to skip the tail split. + torch::Tensor output_tail; +}; + +struct FusedMlaQParams { + // Query tensor for the MLA attention operation. + // Shape: (batch_size, sequence_length, input_size). + // Dtype: float16 or bfloat16. + torch::Tensor q; + + // Output tensor for the fused MLA query operation. + // Shape: (batch_size, sequence_length, head_num, head_size). + // Dtype: same as q, int8, float8_e4m3fn. + torch::Tensor output; + + // Output quantization scales for dynamic per-token quantization. + // Shape: (batch_size, sequence_length, head_num). + // Dtype: float32. + // Only used when quant_mode is "dynamic_per_token". + torch::Tensor output_scale; + + // Intermediate RMSNorm result tensor. + // Shape: (batch_size, sequence_length, input_size). + // Dtype: same as q. + std::optional output_norm; + + // Scaling parameter for RMSNorm normalization. + // Shape: (input_size). + // Dtype: same as q. + torch::Tensor gamma; + + // Smooth quantization scale for input tensor. + // Shape: (input_size) if provided. + // Dtype: float32. + // Optional: can be nullopt if smooth quantization is not used. + std::optional smooth_quant_scale; + + // Weight matrix for the first matmul operation in MLA. + // Shape: (head_num * (nope_dim + pe_dim), input_size). + // Dtype: int8, float8_e4m3fn. + torch::Tensor weight_b; + + // Per-channel scale for weight_b quantization. + // Shape: (head_num * (nope_dim + pe_dim)). + // Dtype: float32. + torch::Tensor weight_b_scale; + + // Weight matrix for the bmm operation in MLA. + // Shape: (head_num, kv_lora_rank, nope_dim). + // Dtype: same as q. + torch::Tensor weight_c; + + // Sine values for rotary position embedding. + // Shape: (rotary_sequence_length, pe_dim). + // Dtype: same as q. + torch::Tensor sin; + + // Cosine values for rotary position embedding. + // Shape: (rotary_sequence_length, pe_dim). + // Dtype: same as q. + torch::Tensor cos; + + // Position IDs for rotary embedding. + // Shape: (batch_size). + // Dtype: int32. + torch::Tensor position_id; + + // Quantization mode for the operation. + // Supported values: "none", "dynamic_per_token". + // Default: "none". + std::string quant_mode = "none"; + + // Epsilon value for RMSNorm numerical stability. + double eps = 1e-6; + + // Rotary embedding mode flag. + // If true, apply cross rotary embedding (interleaved). + // If false, apply fold rotary embedding (non-interleaved). + bool interleaved = true; +}; + +struct FusedMlaKVParams { + // The input key-value tensor. + // Shape: (batch, seq, head_num, head_size). + // Dtype: half, bfloat16. + torch::Tensor input_kv; + + // The rotary sin table tensor. + // Shape: (rotary_seq, rotary_dim). + // Dtype: same as input_kv. + torch::Tensor sin; + + // The rotary cos table tensor. + // Shape: (rotary_seq, rotary_dim). + // Dtype: same as input_kv. + torch::Tensor cos; + + // The rotary seq_len offset of each batch. + // Shape: (batch). + // Dtype: int32. + torch::Tensor position_id; + + // The weight of RMSNorm normalization. + // Shape: (norm_dim). + // Dtype: same as input_kv. + torch::Tensor gamma; + + // The cache tensor for key-value storage. + // Shape: (num_blocks, num_heads, block_size, head_size). + // Dtype: half, bfloat16, int8, float8_e4m3fn. + torch::Tensor kv_cache; + + // Scale tensor for cache quantization. + // For static per-channel quantization: shape is (head_num, head_size) or + // (batch, head_num, head_size). For dynamic per-token quantization: shape is + // (num_blocks, head_num, block_size) and is an output tensor. Dtype: float32. + // Optional: only used when quant_mode is "static_per_channel" or + // "dynamic_per_token". + std::optional kv_cache_scale; + + // The slot mapping tensor for paged attention. + // Shape: (batch, seq). + // Dtype: int32. + // Optional: only required when is_paged_cache is true. + std::optional slot_mapping; + + // The batch index in the cache where the kv tensors will be placed. + // Shape: (batch). + // Dtype: int32. + // Optional: used for non-paged cache style. + std::optional cache_bs_id; + + // A 1D tensor representing the sequence offsets where the cache data starts + // for each batch. Shape: (batch). Dtype: int32. Optional: used for non-paged + // cache style. + std::optional cache_seq_offset; + + // Quantization mode for the operation. + // Supported values: "none", "static_per_channel", "dynamic_per_token". + std::string quant_mode = "none"; + + // Flag indicating the cache style. + // If true, uses paged cache style and slot_mapping must be provided. + // If false, uses linear cache style and cache_bs_id/cache_seq_offset may be + // used. Default: true. + bool is_paged_cache = true; + + // Epsilon value for RMSNorm numerical stability. + double eps = 1e-6; + + // Rotary embedding mode flag. + // If true, apply cross rotary embedding (interleaved). + // If false, apply fold rotary embedding (non-interleaved). + bool interleaved = true; +}; + +struct FusedIndexerQParams { + // The input tensor for query projection. + // Shape: (token_num, input_dim). + // Dtype: half, bfloat16. + torch::Tensor input_q; + + // An output tensor to store the final result in-place. + // Shape: (token_num, head_num, head_size). + // Dtype: same as input_q, or int8 if output is quantized. + torch::Tensor output; + + // Optional output tensor to store quantization scales. + // Shape: (token_num, head_num). + // Dtype: float32. + std::optional output_scale; + + // The weight tensor for query projection. + // Shape: (head_num, head_size, input_dim). + // Dtype: half, bfloat16. + torch::Tensor w_q; + + // The scale tensor for the w_q weight, used for per-channel quantization. + // Shape: (head_num, head_size). + // Dtype: float32. + std::optional w_q_scale; + + // Optional weight tensor for the Hadamard transformation. + // Shape: (head_size, head_size). + // Dtype: same as input_q. + std::optional hadamard_matrix; + + // A pre-computed tensor containing sine values for RoPE. + // Shape: (rotary_seq, rotary_dim). + // Dtype: same as input_q. + torch::Tensor sin; + + // A pre-computed tensor containing cosine values for RoPE. + // Shape: (rotary_seq, rotary_dim). + // Dtype: same as input_q. + torch::Tensor cos; + + // A tensor indicating the position index for each token. + // Shape: (token_num). + // Dtype: int32. + torch::Tensor position_id; + + // Quantization mode for the output. + // Supported values: "none", "dynamic_per_token". + std::string quant_mode = "none"; + + // Rotary embedding mode flag. + // If true, apply cross rotary embedding (interleaved). + // If false, apply fold rotary embedding (non-interleaved). + bool interleaved = true; + + // Flag indicating whether to apply RoPE at the front of the operation. + // If true, apply RoPE at the front of the operation. + // If false, apply RoPE at the back of the operation. + bool rope_at_front = true; +}; + +struct FusedIndexerKParams { + // The input tensor. + // Shape: (m, dim). + // Dtype: half, bfloat16. + torch::Tensor x; + + // The weight tensor for K projection. + // Shape: (head_size, dim). + // Dtype: same as x. + torch::Tensor wk; + + // The weight tensor for head projection. + // Shape: (head_num, dim). + // Dtype: same as x. + torch::Tensor wproj; + + // A pre-computed tensor containing sine values for RoPE. + // Shape: (rotary_seq, rope_dim). + // Dtype: same as x. + torch::Tensor sin_table; + + // A pre-computed tensor containing cosine values for RoPE. + // Shape: (rotary_seq, rope_dim). + // Dtype: same as x. + torch::Tensor cos_table; + + // A tensor indicating the position index for each token. + // Shape: (m). + // Dtype: int32. + torch::Tensor position_id; + + // A tensor mapping tokens to cache slots. + // Shape: (m). + // Dtype: int32. + torch::Tensor slot_mapping; + + // The computed head weights tensor. + // Shape: (m, head_num). + // Dtype: same as x. + torch::Tensor head_weights; + + // The K cache tensor. + // Shape: (block_num, 1, block_size, head_size). + // Dtype: half, bfloat16, int8. + torch::Tensor k_cache; + + // Optional scale tensor for quantized K cache. + // Shape: (block_num, 1, block_size). + // Dtype: float32. + std::optional k_cache_scale; + + // Optional weight tensor for the Hadamard transformation. + // Shape: (head_size, head_size). + // Dtype: same as x. + std::optional hadamard_matrix; + + // Rotary embedding mode flag. + // If true, apply cross rotary embedding (interleaved). + // If false, apply fold rotary embedding (non-interleaved). + bool interleaved = true; + + // Optional weight tensor for RMSNorm. + // Shape: (head_size). + // Dtype: float32. + std::optional gamma; + + // Optional bias tensor for RMSNorm. + // Shape: (head_size). + // Dtype: float32. + std::optional beta; + + // RMSNorm epsilon. + double eps = 1e-6; +}; + +struct MoeInitRoutingV2Params { + // TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx + + // moe_expand_input (and token_count/cusum outputs) on other backends. + torch::Tensor x; + torch::Tensor expert_idx; + std::optional scale; + std::optional offset; + int active_num; + int expert_capacity; + int expert_num; + int drop_pad_mode; + int expert_tokens_num_type; + bool expert_tokens_num_flag; + int quant_mode; + torch::IntArrayRef active_expert_range; + int row_idx_type; +}; + +// FP8 scaled quantize parameters +// Quantizes input tensor to FP8 e4m3 format with scale +struct Fp8ScaledQuantizeParams { + // Input tensor. Shape: [M, K]. Dtype: float16, bfloat16. + torch::Tensor input; + // Optional output tensor. Shape: [M, K]. Dtype: float8_e4m3fn. + // If not provided, will be allocated automatically. + std::optional output; + // Optional pre-computed scale for static quantization. + // Shape: scalar or [1]. If not provided, scale will be computed dynamically. + std::optional scale; +}; + +// FP8 scaled matmul parameters for W8A8 quantization +// Performs: c = (a @ b.T) with scales applied, following CUTLASS convention +struct Fp8ScaledMatmulParams { + // Quantized input tensor A. Shape: [M, K]. Dtype: float8_e4m3fn. + torch::Tensor a; + // Quantized weight tensor B. Shape: [N, K] (will be transposed internally). + // Dtype: float8_e4m3fn. + torch::Tensor b; + // Scale for tensor A. Shape: scalar or [1]. + torch::Tensor a_scale; + // Scale for tensor B. Shape: scalar or [1]. + torch::Tensor b_scale; + // Optional bias tensor. Shape: [N]. + std::optional bias; + // Optional output tensor. Shape: [M, N]. + // If not provided, will be allocated with output_dtype. + std::optional output; + // Output data type. Typically float16 or bfloat16. + torch::ScalarType output_dtype; + // Optional original input shape (before flatten to 2D). + // If provided, output will be reshaped to match original input dimensions. + // E.g., input_shape = [batch, seq, hidden] -> output = [batch, seq, N] + std::optional> input_shape; +}; + +// Static scaled FP8 quantization parameters +// Quantizes input tensor to FP8 using a pre-computed scale factor +struct StaticScaledFp8QuantParams { + // Output tensor to store quantized result. Shape: [..., d]. + // Dtype: float8_e4m3fn. Must be pre-allocated. + torch::Tensor output; + // Input tensor to quantize. Shape: [..., d]. + // Dtype: float16, bfloat16, or float32. + torch::Tensor input; + // Pre-computed scale factor. Shape: [1] or scalar. + // Dtype: float32. Used for static quantization. + torch::Tensor scale; +}; + +// Fused RMSNorm + Static FP8 Quantization Parameters +// These fused operations combine RMSNorm and FP8 quantization to reduce memory +// bandwidth by avoiding the intermediate write-back to global memory. + +// Fused RMSNorm + Static FP8 Quantization parameters (without residual) +struct RmsNormStaticFp8QuantParams { + // Input tensor. Shape: [..., hidden_size]. Dtype: float16, bfloat16, float32. + torch::Tensor input; + // RMSNorm weight. Shape: [hidden_size]. Dtype: same as input. + torch::Tensor weight; + // FP8 quantization scale (pre-computed). Shape: [1]. Dtype: float32. + torch::Tensor scale; + // RMSNorm epsilon. + double epsilon; +}; + +// Fused Add + RMSNorm + Static FP8 Quantization parameters (with residual) +struct FusedAddRmsNormStaticFp8QuantParams { + // Input tensor. Shape: [..., hidden_size]. Dtype: float16, bfloat16, float32. + torch::Tensor input; + // Residual tensor. Shape: [..., hidden_size]. Dtype: same as input. + // Updated in-place with: residual = input + residual + torch::Tensor residual; + // RMSNorm weight. Shape: [hidden_size]. Dtype: same as input. + torch::Tensor weight; + // FP8 quantization scale (pre-computed). Shape: [1]. Dtype: float32. + torch::Tensor scale; + // RMSNorm epsilon. + double epsilon; +}; + +// NPU Fused GDN Gating parameters +struct FusedGdnGatingParams { + torch::Tensor A_log; + torch::Tensor a; + torch::Tensor b; + torch::Tensor dt_bias; + float beta = 1.0f; + float threshold = 20.0f; +}; + +// NPU Fused Recurrent Gated Delta Rule parameters +struct FusedRecurrentGatedDeltaRuleParams { + torch::Tensor q; + torch::Tensor k; + torch::Tensor v; + torch::Tensor g; + std::optional beta = std::nullopt; + std::optional scale = std::nullopt; + std::optional initial_state = std::nullopt; + bool inplace_final_state = true; + std::optional cu_seqlens = std::nullopt; + std::optional ssm_state_indices = std::nullopt; + std::optional num_accepted_tokens = std::nullopt; + bool use_qk_l2norm_in_kernel = false; +}; + +// NPU Causal Conv1d Update parameters +struct CausalConv1dUpdateParams { + torch::Tensor x; + torch::Tensor conv_state; + torch::Tensor weight; + bool activation = true; + std::optional bias = std::nullopt; + std::optional conv_state_indices = std::nullopt; + std::optional query_start_loc = std::nullopt; + int32_t max_query_len = -1; + int32_t pad_slot_id = -1; + std::optional block_idx_last_scheduled_token; + std::optional initial_state_idx; + bool validate_data = false; +}; + +struct GatedLayerNormParams { + torch::Tensor x; + torch::Tensor weight; + torch::Tensor bias; + double eps; + std::optional z = std::nullopt; + int64_t group_size = -1; + bool norm_before_gate = true; + bool is_rms_norm = true; +}; + +struct PartialRotaryEmbeddingParams { + torch::Tensor positions; + torch::Tensor query; + torch::Tensor key; + int64_t head_size; + int64_t rotary_dim; + torch::Tensor cos_sin_cache; + bool is_neox_style; +}; + +struct FusedQkvzbaSplitReshapeParams { + torch::Tensor mixed_qkvz; + torch::Tensor mixed_ba; + int32_t num_heads_qk; + int32_t num_heads_v; + int32_t head_qk; + int32_t head_v; +}; + +struct GemmaRMSNormParams { + torch::Tensor x; + torch::Tensor gamma; + double epsilon; + torch::Tensor rstd_out; + torch::Tensor norm_out; +}; + +struct SplitQkvRmsnormMropeParams { + torch::Tensor qkvg; + torch::Tensor q_weight; + torch::Tensor k_weight; + torch::Tensor cos_sin; + torch::Tensor gather_pattern; + float eps; + int64_t num_q_heads; + int64_t num_kv_heads; + int64_t head_size; +}; + +struct ChunkGatedDeltaRuleParams { + // Query tensor. Shape: [B, T, Hqk, K]. Dtype: bfloat16. + torch::Tensor q; + // Key tensor. Shape: [B, T, Hqk, K]. Dtype: bfloat16. + torch::Tensor k; + // Value tensor. Shape: [B, T, H, V]. Dtype: bfloat16. + torch::Tensor v; + // Gating tensor. Shape: [B, T, H]. Dtype: float32 or bfloat16. + torch::Tensor g; + // Beta tensor. Shape: [B, T, H]. Dtype: float32 or bfloat16. + torch::Tensor beta; + // Optional scale factor for attention. Default: K^(-0.5). + std::optional scale = std::nullopt; + // Optional initial state tensor. Shape: [N, H, K, V]. Dtype: bfloat16. + std::optional initial_state = std::nullopt; + // Whether to output the final state. + bool output_final_state = false; + // Chunk size for processing. Default: 64. + int64_t chunk_size = 64; + // Optional cumulative sequence lengths. Shape: [num_sequences + 1]. Dtype: + // int32. + std::optional cu_seqlens = std::nullopt; + // Whether input is head-first format. Default: false (batch-first). + bool head_first = false; + // Whether to apply L2 norm to q and k inside the kernel. Default: false. + bool use_qk_l2norm_in_kernel = false; +}; +} // namespace xllm::kernel