diff --git a/README.md b/README.md index 37dcf9dbc..d803b3d5d 100644 --- a/README.md +++ b/README.md @@ -195,3 +195,8 @@ LM 不隐式构建依赖栈，也不安装 Core 的独立 Python runtime。 ```bash python examples/bench.py --device nvidia --model= --enable-paged-attn [--attn=default | --attn=flash-attn] ``` + +### Clover on Hygon + +Native Clover block diffusion and the optional frontier/expert-reuse decoder are +documented in [the Hygon guide](docs/clover-hygon.md). diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index b2243ef25..b0a0e970e 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -448,6 +448,8 @@ InferEngine::Input::to_model_input( : to_device(target_hidden_states), sample_all_positions}; input.logical_kv_upper_bound = logical_kv_upper_bound; + input.diffusion_preview = diffusion_preview; + input.diffusion_expert_budget = diffusion_expert_budget; auto &attn_metadata = forward_context.attn_metadata; attn_metadata = { diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index 4fcbb3e0b..d8193bbb9 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -98,6 +98,8 @@ class RankWorker { std::optional>> sampling_rows; std::optional nsa_owner_snapshot; + bool diffusion_preview{false}; + bool diffusion_expert_budget{false}; infinilm::InfinilmModel::Input to_model_input( infinilm::runtime::Device device, diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index 76e7a3939..f47a0e9e4 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -31,6 +31,7 @@ struct StaticPrefillHostDescription { }; struct AttentionMetadata { + bool diffusion_preview{false}; /// Past Lengths of cached sequence for each request, of shape `[num_requests]`. std::optional past_sequence_lengths; /// ToTal Lengths for each request sequence, of shape `[num_requests]`. @@ -119,6 +120,11 @@ struct ForwardContext { AttentionMetadata attn_metadata; MambaMetadata mamba_metadata; MultiModalMetadata mm_metadata; + bool diffusion_preview{false}; + size_t diffusion_block_size{32}; + bool diffusion_expert_budget{false}; + infinilm::runtime::Tensor diffusion_valid; + infinilm::runtime::Tensor diffusion_live; std::vector kv_cache_vec; std::vector conv_state_vec; std::vector ssm_state_vec; diff --git a/csrc/layers/attention/backends/static_attn.cpp b/csrc/layers/attention/backends/static_attn.cpp index 31e5bb758..6f18b3e37 100644 --- a/csrc/layers/attention/backends/static_attn.cpp +++ b/csrc/layers/attention/backends/static_attn.cpp @@ -95,6 +95,7 @@ infinilm::runtime::Tensor StaticAttentionImpl::forward(const AttentionLayer &lay const bool finite_window = window_size[0] >= 0; bool varlen_candidate = finite_window; varlen_candidate = varlen_candidate || (prefill && seq_len > 1 && kv_quant_scheme_ == infinilm::quantization::KVQuantAlgo::NONE && (q_rope->dtype() == infinilm::runtime::DataType::kFloat16 || q_rope->dtype() == infinilm::runtime::DataType::kBFloat16)); + varlen_candidate = varlen_candidate && !attn_metadata.diffusion_preview; if (finite_window && seq_len > 1 && !prefill) { throw std::invalid_argument("Static finite window requires validated prefill row/past/total metadata"); } @@ -269,7 +270,7 @@ infinilm::runtime::Tensor StaticAttentionImpl::forward(const AttentionLayer &lay throw std::invalid_argument("causal_softmax has unsupported shape or dtype"); } std::optional causal_mask; - if (seq_len != 1) { + if (seq_len != 1 && !attn_metadata.diffusion_preview) { causal_mask = infinilm::runtime::Tensor::empty( runtime, {seq_len, total_seq_len}, attn_weight_softmax->dtype()); } diff --git a/csrc/models/clover_moe/clover_decoder_layer.cpp b/csrc/models/clover_moe/clover_decoder_layer.cpp new file mode 100644 index 000000000..0e8f3c2af --- /dev/null +++ b/csrc/models/clover_moe/clover_decoder_layer.cpp @@ -0,0 +1,83 @@ +#include "clover_decoder_layer.hpp" + +#include "runtime/runtime.hpp" +#include +#include +#include +#include + +namespace infinilm::models::clover_moe { + +CloverDecoderLayer::CloverDecoderLayer(infinilm::runtime::Runtime &runtime, + std::shared_ptr model_config, + size_t layer_idx, + const engine::distributed::RankInfo &rank_info, + backends::AttentionBackend attention_backend, + const layers::moe::EPConfig &ep_config) + : layer_idx_(layer_idx) { + runtime.assertOwnerThread(); + if (runtime.isGraphRecording()) { + throw std::logic_error("Clover decoder construction requires eager Runtime"); + } + if (!model_config || rank_info.device != runtime.device() || rank_info.tp_size <= 0 + || rank_info.tp_rank < 0 || rank_info.tp_rank >= rank_info.tp_size + || (rank_info.tp_size > 1 && !rank_info.comm)) { + throw std::invalid_argument("Invalid explicit Qwen35 decoder construction resources"); + } + const auto layer_types = model_config->get>("layer_types"); + layer_type_ = layer_types.at(layer_idx); + if (layer_type_ != "linear_attention" && layer_type_ != "full_attention") { + throw std::invalid_argument("Clover decoder has unsupported layer_type"); + } + try { + const auto &dtype = model_config->get_dtype(); + const size_t hidden_size = model_config->get("hidden_size"); + const double rms_norm_eps = model_config->get("rms_norm_eps"); + + input_layernorm_ = this->bind_child("input_layernorm", std::make_shared>(runtime, hidden_size, rms_norm_eps, dtype)); + post_attention_layernorm_ = this->bind_child("post_attention_layernorm", std::make_shared>(runtime, hidden_size, rms_norm_eps, dtype)); + mlp_ = this->bind_child("mlp", std::make_shared>(runtime, model_config, layer_idx, rank_info, attention_backend, ep_config)); + + if ("linear_attention" == layer_type_) { + linear_attn_ = this->bind_child("linear_attn", std::make_shared>(runtime, model_config, layer_idx, rank_info, attention_backend, ep_config)); + } else if ("full_attention" == layer_type_) { + self_attn_ = this->bind_child("self_attn", std::make_shared>(runtime, model_config, layer_idx, rank_info, attention_backend)); + } else { + throw std::runtime_error( + "CloverDecoderLayer: unsupported layer_type '" + layer_type_ + + "' for layer " + std::to_string(layer_idx)); + } + } catch (...) { + try { + runtime.syncStream(); + } catch (...) { std::terminate(); } + throw; + } +} + +std::tuple CloverDecoderLayer::forward( + const infinilm::runtime::Tensor &positions, + infinilm::runtime::Tensor &hidden_states, + infinilm::runtime::Tensor &residual, + infinilm::global_state::ForwardContext &forward_context) { + auto &runtime = forward_context.require_runtime(); + runtime.assertOwnerThread(); + if (!hidden_states || hidden_states->device() != runtime.device() + || (residual && residual->device() != runtime.device()) + || input_layernorm_->weight()->device() != runtime.device() + || post_attention_layernorm_->weight()->device() != runtime.device()) { + throw std::invalid_argument("CloverDecoderLayer: forward Runtime/input/norm device mismatch"); + } + input_layernorm_->forward_inplace(runtime, hidden_states, residual); + if ("linear_attention" == layer_type_) { + hidden_states = linear_attn_->forward(hidden_states, forward_context); + } else { + hidden_states = self_attn_->forward(positions, hidden_states, forward_context); + } + + post_attention_layernorm_->forward_inplace(runtime, hidden_states, residual); + hidden_states = mlp_->forward(runtime, hidden_states, forward_context); + return std::make_tuple(hidden_states, residual); +} + +} // namespace infinilm::models::clover_moe diff --git a/csrc/models/clover_moe/clover_decoder_layer.hpp b/csrc/models/clover_moe/clover_decoder_layer.hpp new file mode 100644 index 000000000..cbe7503fa --- /dev/null +++ b/csrc/models/clover_moe/clover_decoder_layer.hpp @@ -0,0 +1,43 @@ +#pragma once +#include "../../layers/normalization/rmsnorm.hpp" +#include "../../weights/weights.hpp" + +#include "../qwen3_5/qwen3_5_attention.hpp" +#include "clover_gdn.hpp" +#include "clover_moe.hpp" + +#include +#include + +namespace infinilm::models::clover_moe { + +class CloverDecoderLayer : public infinilm::weights::ModelWeights { +public: + CloverDecoderLayer(infinilm::runtime::Runtime &runtime, + std::shared_ptr model_config, + size_t layer_idx, + const engine::distributed::RankInfo &rank_info, + backends::AttentionBackend attention_backend, + const layers::moe::EPConfig &ep_config); + + std::tuple forward( + const infinilm::runtime::Tensor &positions, + infinilm::runtime::Tensor &hidden_states, + infinilm::runtime::Tensor &residual, + infinilm::global_state::ForwardContext &forward_context); + + size_t layer_idx() const { return layer_idx_; } + +protected: + std::shared_ptr input_layernorm_; + std::shared_ptr post_attention_layernorm_; + std::shared_ptr self_attn_; + std::shared_ptr linear_attn_; + std::shared_ptr mlp_; + +private: + size_t layer_idx_; + std::string layer_type_; +}; + +} // namespace infinilm::models::clover_moe diff --git a/csrc/models/clover_moe/clover_for_causal_lm.cpp b/csrc/models/clover_moe/clover_for_causal_lm.cpp new file mode 100644 index 000000000..f38131fa6 --- /dev/null +++ b/csrc/models/clover_moe/clover_for_causal_lm.cpp @@ -0,0 +1,105 @@ +#include "../../layers/causal_lm_templates/text_causal_lm.hpp" +#include "../models_registry.hpp" +#include "../qwen3_5/qwen3_5_for_causal_lm.hpp" +#include "../qwen3_next/qwen3_next_allocate_kv_cache_tensors.hpp" +#include "clover_decoder_layer.hpp" +#include + +namespace infinilm::models::clover_moe { + +using CloverModel = layers::causal_lm_templates::TextModel; + +class CloverForCausalLM : public layers::causal_lm_templates::TextCausalLM { +public: + using Base = layers::causal_lm_templates::TextCausalLM; + using Base::Base; + + Output forward(const Input &input, global_state::ForwardContext &context) const override { + auto &runtime = context.require_runtime(); + if (runtime.isGraphRecording()) { + throw std::invalid_argument("Clover requires eager execution."); + } + context.diffusion_preview = input.diffusion_preview; + context.diffusion_expert_budget = input.diffusion_expert_budget; + context.attn_metadata.diffusion_preview = input.diffusion_preview; + context.diffusion_block_size = model_config_->get("block_size"); + if (!input.diffusion_preview) { + denoise_step_ = 0; + } + if (input.diffusion_preview) { + const auto tokens = input.input_ids.value()->numel(); + if (!input.sample_all_positions || tokens % context.diffusion_block_size) { + throw std::invalid_argument("Clover preview requires complete blocks and all-position logits."); + } + context.diffusion_valid = runtime::Tensor::empty(runtime, {tokens}, runtime::DataType::kInt8); + execution::submit(runtime, infini::ops::Fill::DefaultConfig(runtime.device().type()), + context.diffusion_valid, 1.0, context.diffusion_valid); + if (input.diffusion_expert_budget) { + bool refresh = denoise_step_ % 5 == 0 || !previous_budget_; + const auto &ids = input.input_ids.value(); + if (!previous_ids_ || previous_ids_->shape() != ids->shape()) { + previous_ids_ = runtime::Tensor::empty(runtime, ids->shape(), ids->dtype()); + refresh = true; + } + context.diffusion_live = runtime::Tensor::empty(runtime, {tokens}, runtime::DataType::kInt8); + execution::submit(runtime, execution::RecordingPolicy::HostOnly, + infini::ops::CloverLiveMask::DefaultConfig(runtime.device().type()), ids, previous_ids_, + static_cast(model_config_->get("mask_token_id")), refresh, context.diffusion_live); + } + previous_budget_ = input.diffusion_expert_budget; + ++denoise_step_; + } + return Base::forward(input, context); + } + + void reset_cache(const cache::CacheConfig *cache_config, global_state::ForwardContext &context, + backends::AttentionBackend backend, const engine::distributed::RankInfo &rank) override { + if (backend != backends::AttentionBackend::STATIC_ATTN) { + throw std::invalid_argument("Clover currently requires the static attention backend."); + } + auto &runtime = context.require_runtime(); + denoise_step_ = 0; + previous_budget_ = false; + previous_ids_.reset(); + auto allocated = qwen3_next::qwen3_next_allocate_cache_tensors(runtime, rank, cache_config, model_config_, backend); + for (auto &state : allocated.ssm_state_tensors) { + if (!state) { + continue; + } + state = runtime::Tensor::empty(runtime, state->shape(), runtime::DataType::kFloat32); + execution::submit(runtime, infini::ops::Fill::DefaultConfig(runtime.device().type()), state, 0.0, state); + } + context.kv_cache_vec = std::move(allocated.kv_cache_tensors); + context.conv_state_vec = std::move(allocated.conv_state_tensors); + context.ssm_state_vec = std::move(allocated.ssm_state_tensors); + cache_config_ = cache_config ? cache_config->unique_copy() : nullptr; + } + +private: + mutable size_t denoise_step_{0}; + mutable bool previous_budget_{false}; + mutable runtime::Tensor previous_ids_; +}; + +std::shared_ptr create_clover_config(std::shared_ptr config) { + if (config->get_dtype() != runtime::DataType::kBFloat16) { + throw std::invalid_argument("Clover requires BF16 model weights."); + } + auto &json = config->get_config_json(); + if (json.at("model_type") != "clover_moe" || json.at("block_size") != 32 || json.at("linear_key_head_dim") != 128 || json.at("linear_value_head_dim") != 128 || json.at("linear_conv_kernel_dim") != 4 || json.at("num_experts") != 256 || json.at("num_experts_per_tok") != 8) { + throw std::invalid_argument("Unsupported Clover architecture."); + } + json["norm_topk_prob"] = true; + json["position_id_axes"] = 1; + json["rope_parameters"]["mrope_section"] = {16, 8, 8}; + json["rope_parameters"]["mrope_interleaved"] = false; + return qwen3_5::prepare_qwen3_5_model_config(config); +} + +} // namespace infinilm::models::clover_moe + +namespace { +INFINILM_REGISTER_EXPLICIT_CAUSAL_LM_MODEL(clover_moe, + infinilm::models::clover_moe::CloverForCausalLM, + infinilm::models::clover_moe::create_clover_config); +} diff --git a/csrc/models/clover_moe/clover_gdn.cpp b/csrc/models/clover_moe/clover_gdn.cpp new file mode 100644 index 000000000..85dde42b7 --- /dev/null +++ b/csrc/models/clover_moe/clover_gdn.cpp @@ -0,0 +1,104 @@ +#include "clover_gdn.hpp" +#include "../../global_state/forward_context.hpp" +#include "execution/submit.hpp" +#include "runtime/runtime.hpp" +#include +#include +#include + +namespace infinilm::models::clover_moe { + +CloverGatedDeltaNet::CloverGatedDeltaNet(runtime::Runtime &runtime, std::shared_ptr config, + size_t layer, const engine::distributed::RankInfo &rank, + backends::AttentionBackend backend, const layers::moe::EPConfig &ep) + : Qwen3NextGatedDeltaNet(runtime, config, layer, rank, backend, ep), + tp_rank_(rank.tp_rank), tp_size_(rank.tp_size), + full_key_dim_(local_key_dim_ * tp_size_), full_value_dim_(local_value_dim_ * tp_size_) { + try { + const auto dtype = config->get_dtype(); + right_weight_ = weights::Weight(runtime, {2 * full_key_dim_ + full_value_dim_, 1, 3}, dtype); + bind_weight("conv_right_weight", right_weight_); + const auto heads = config->get("linear_num_value_heads"); + logit_scale_ = weights::Weight(runtime, {heads}, dtype, 0, tp_rank_, tp_size_); + history_bias_ = weights::Weight(runtime, {heads}, dtype, 0, tp_rank_, tp_size_); + history_c_ = weights::Weight(runtime, {heads}, dtype, 0, tp_rank_, tp_size_); + bind_weight("bd_softmax_logit_scale", logit_scale_); + bind_weight("bd_softmax_history_bias", history_bias_); + bind_weight("bd_softmax_history_c", history_c_); + } catch (...) { + runtime.syncStream(); + throw; + } +} + +void CloverGatedDeltaNet::process_weights_after_loading(runtime::Runtime &runtime, bool) { + conv_weight_ = conv1d_->state_dict().at("weight").tensor(); + auto local = runtime::Tensor::empty(runtime, {2 * local_key_dim_ + local_value_dim_, 1, 3}, right_weight_->dtype()); + for (size_t part = 0; part < 3; ++part) { + const size_t width = part == 2 ? local_value_dim_ : local_key_dim_; + const size_t src = part * full_key_dim_ + tp_rank_ * width; + const size_t dst = part * local_key_dim_; + local->narrow({{0, dst, width}})->copy_from(runtime, right_weight_->narrow({{0, src, width}})); + } + runtime.syncStream(); + right_weight_ = weights::Weight(local); + local_weights_["conv_right_weight"] = right_weight_; +} + +runtime::Tensor CloverGatedDeltaNet::forward(const runtime::Tensor &hidden, + const global_state::ForwardContext &context) const { + auto &runtime = context.require_runtime(); + auto input = hidden; + const size_t tokens = hidden->size(0) * hidden->size(1); + auto qkv = in_proj_qkv_->forward(runtime, input)->view({1, tokens, 2 * local_key_dim_ + local_value_dim_}); + auto z = in_proj_z_->forward(runtime, input); + auto a = in_proj_a_->forward(runtime, input)->view({1, tokens, local_num_value_heads_}); + auto b = in_proj_b_->forward(runtime, input)->view({1, tokens, local_num_value_heads_}); + const auto device = runtime.device().type(); + runtime::Tensor mixed; + if (context.diffusion_preview) { + auto flat = qkv->view({tokens, qkv->size(2)}); + auto out = runtime::Tensor::empty(runtime, flat->shape(), flat->dtype()); + execution::submit( + runtime, execution::RecordingPolicy::HostOnly, + infini::ops::CloverBlockConv1d::DefaultConfig(device), + flat, context.conv_state_vec[layer_idx_], context.mamba_metadata.init_state_indices.value(), + conv_weight_, right_weight_.tensor(), context.diffusion_valid, context.diffusion_block_size, out); + mixed = out->view(qkv->shape()); + } else { + mixed = conv1d_->forward(qkv, context); + } + auto q = mixed->narrow({{2, 0, local_key_dim_}})->as_strided({1, tokens, local_num_key_heads_, key_head_dim_}, {mixed->stride(0), mixed->stride(1), static_cast(key_head_dim_), 1}); + auto k = mixed->narrow({{2, local_key_dim_, local_key_dim_}})->as_strided(q->shape(), q->strides()); + auto v = mixed->narrow({{2, 2 * local_key_dim_, local_value_dim_}})->as_strided({1, tokens, local_num_value_heads_, value_head_dim_}, {mixed->stride(0), mixed->stride(1), static_cast(value_head_dim_), 1}); + auto g = runtime::Tensor::empty(runtime, a->shape(), runtime::DataType::kFloat32); + auto beta = runtime::Tensor::empty(runtime, a->shape(), runtime::DataType::kFloat32); + execution::submit( + runtime, execution::RecordingPolicy::HostOnly, + infini::ops::FusedGatedDeltaNetGating::DefaultConfig(device), + A_log_.tensor(), a, b, dt_bias_.tensor(), 1.0f, 20.0f, g, beta); + auto out = runtime::Tensor::empty(runtime, v->shape(), v->dtype()); + if (context.diffusion_preview) { + execution::submit( + runtime, execution::RecordingPolicy::HostOnly, + infini::ops::CloverLocalMemory::DefaultConfig(device), + q, k, v, g, context.ssm_state_vec[layer_idx_], context.mamba_metadata.init_state_indices.value(), + context.diffusion_valid, logit_scale_.tensor(), history_bias_.tensor(), history_c_.tensor(), + context.diffusion_block_size, out); + } else { + execution::submit( + runtime, execution::RecordingPolicy::HostOnly, + infini::ops::CloverRecurrentGdn::DefaultConfig(device), + q, k, v, g, beta, context.mamba_metadata.input_offsets.value(), + context.mamba_metadata.init_state_indices.value(), context.mamba_metadata.final_state_indices.value(), + context.ssm_state_vec[layer_idx_], out); + } + auto normalized = norm_->forward(runtime, out->view({tokens * local_num_value_heads_, value_head_dim_})); + normalized = normalized->view(z->shape()); + auto gate = runtime::Tensor::empty(runtime, z->shape(), z->dtype()); + execution::submit(runtime, infini::ops::Silu::DefaultConfig(device), z, gate); + execution::submit(runtime, infini::ops::Mul::DefaultConfig(device), normalized, gate, normalized); + return out_proj_->forward(runtime, normalized); +} + +} // namespace infinilm::models::clover_moe diff --git a/csrc/models/clover_moe/clover_gdn.hpp b/csrc/models/clover_moe/clover_gdn.hpp new file mode 100644 index 000000000..36589dc1d --- /dev/null +++ b/csrc/models/clover_moe/clover_gdn.hpp @@ -0,0 +1,22 @@ +#pragma once + +#include "../qwen3_next/qwen3_next_gated_deltanet.hpp" + +namespace infinilm::models::clover_moe { + +class CloverGatedDeltaNet : public qwen3_next::Qwen3NextGatedDeltaNet { +public: + CloverGatedDeltaNet(runtime::Runtime &runtime, std::shared_ptr config, + size_t layer, const engine::distributed::RankInfo &rank, + backends::AttentionBackend backend, const layers::moe::EPConfig &ep); + runtime::Tensor forward(const runtime::Tensor &hidden, + const global_state::ForwardContext &context) const; + void process_weights_after_loading(runtime::Runtime &runtime, bool pre_transpose) override; + +private: + weights::Weight right_weight_, logit_scale_, history_bias_, history_c_; + runtime::Tensor conv_weight_; + size_t tp_rank_, tp_size_, full_key_dim_, full_value_dim_; +}; + +} // namespace infinilm::models::clover_moe diff --git a/csrc/models/clover_moe/clover_moe.cpp b/csrc/models/clover_moe/clover_moe.cpp new file mode 100644 index 000000000..9b699680c --- /dev/null +++ b/csrc/models/clover_moe/clover_moe.cpp @@ -0,0 +1,104 @@ +#include "clover_moe.hpp" +#include "../../global_state/forward_context.hpp" +#include "../../layers/linear/dense.hpp" +#include "execution/submit.hpp" +#include +#include +#include + +namespace infinilm::models::clover_moe { + +CloverSparseMoe::CloverSparseMoe(runtime::Runtime &runtime, std::shared_ptr config, + size_t layer, const engine::distributed::RankInfo &rank, + backends::AttentionBackend backend, const layers::moe::EPConfig &ep) + : Qwen3NextSparseMoeBlock(runtime, config, layer, rank, backend, ep), + coreset_(config->get_or("moe_block_coreset_enabled", true)), + top_p_(config->get_or("moe_block_coreset_top_p", 0.9)), + top_k_(config->get("num_experts_per_tok")) { + router_weight_ = gate_->state_dict().at("weight").tensor(); +} + +std::tuple CloverSparseMoe::route( + runtime::Runtime &runtime, const runtime::Tensor &hidden, + const global_state::ForwardContext &context) const { + if (!context.diffusion_preview || (!coreset_ && !context.diffusion_expert_budget)) { + return gate_->forward(runtime, hidden); + } + auto logits = layers::linear::dense_projection(runtime, hidden, router_weight_, std::nullopt, 1.0f); + auto masked = logits; + const auto device = runtime.device().type(); + if (coreset_) { + masked = runtime::Tensor::empty(runtime, logits->shape(), logits->dtype()); + execution::submit(runtime, execution::RecordingPolicy::HostOnly, + infini::ops::CloverCoreset::DefaultConfig(device), logits, context.diffusion_block_size, top_p_, masked); + } + const auto active_top_k = context.diffusion_expert_budget ? size_t{4} : top_k_; + auto weights = runtime::Tensor::empty(runtime, {hidden->size(0), active_top_k}, runtime::DataType::kFloat32); + auto indices = runtime::Tensor::empty(runtime, weights->shape(), runtime::DataType::kInt32); + auto scratch = runtime::Tensor::empty(runtime, weights->shape(), runtime::DataType::kInt32); + execution::submit(runtime, execution::RecordingPolicy::HostOnly, + infini::ops::TopkSoftmax::DefaultConfig(device), masked, + std::optional{}, std::optional{}, true, + weights, indices, scratch); + return {weights, indices}; +} + +infinilm::runtime::Tensor CloverSparseMoe::forward(infinilm::runtime::Runtime &runtime, + const infinilm::runtime::Tensor &hidden_states, + const global_state::ForwardContext &context) const { + ASSERT(hidden_states->ndim() == 3); + + auto shape = hidden_states->shape(); + auto hidden_states_reshaped = hidden_states->view({shape[0] * shape[1], shape[2]}); + + auto [routing_weights, selected_experts] = route(runtime, hidden_states_reshaped, context); + if (context.diffusion_preview && context.diffusion_expert_budget) { + execution::submit(runtime, execution::RecordingPolicy::HostOnly, + infini::ops::CloverMaskRoutes::DefaultConfig(runtime.device().type()), + selected_experts, context.diffusion_live, int64_t{256}, selected_experts); + } + infinilm::layers::moe::TopKOutput topk_output{ + routing_weights, + selected_experts, + infinilm::runtime::Tensor(), + }; + auto routed_states = fused_moe_->forward( + runtime, hidden_states_reshaped, + topk_output, + experts_->moe_weights()); + if (context.diffusion_preview && context.diffusion_expert_budget) { + if (!cached_routed_ || cached_routed_->shape() != routed_states->shape()) { + cached_routed_ = runtime::Tensor::empty(runtime, routed_states->shape(), routed_states->dtype()); + } + execution::submit(runtime, execution::RecordingPolicy::HostOnly, + infini::ops::CloverMergeMoe::DefaultConfig(runtime.device().type()), + routed_states, context.diffusion_live, cached_routed_, routed_states); + } + + auto shared_states = shared_expert_->forward(runtime, hidden_states); + auto hidden_states_for_gate = hidden_states; + const auto ops_input = shared_expert_gate_->forward(runtime, hidden_states_for_gate); + auto shared_gate = infinilm::runtime::Tensor::empty(runtime, ops_input->shape(), ops_input->dtype()); + infinilm::execution::submit( + runtime, infini::ops::Sigmoid::DefaultConfig(runtime.device().type()), ops_input, shared_gate); + INFINILM_TRACE_OP(Sigmoid, shared_gate); + shared_gate = shared_gate->as_strided(shared_states->shape(), {shared_gate->stride(0), shared_gate->stride(1), 0}); + { + auto ops_output = infinilm::runtime::Tensor::empty(runtime, shared_states->shape(), shared_states->dtype()); + infinilm::execution::submit( + runtime, infini::ops::Mul::DefaultConfig(runtime.device().type()), shared_states, shared_gate, ops_output); + INFINILM_TRACE_OP(Mul, ops_output); + shared_states = std::move(ops_output); + } + + auto routed_states_3d = routed_states->as_strided( + {shape[0], shape[1], shape[2]}, + {static_cast(shape[1] * shape[2]), static_cast(shape[2]), 1}); + auto ops_output = infinilm::runtime::Tensor::empty(runtime, routed_states_3d->shape(), routed_states_3d->dtype()); + infinilm::execution::submit( + runtime, infini::ops::Add::DefaultConfig(runtime.device().type()), routed_states_3d, shared_states, ops_output); + INFINILM_TRACE_OP(Add, ops_output); + return ops_output; +} + +} // namespace infinilm::models::clover_moe diff --git a/csrc/models/clover_moe/clover_moe.hpp b/csrc/models/clover_moe/clover_moe.hpp new file mode 100644 index 000000000..8c1bd558e --- /dev/null +++ b/csrc/models/clover_moe/clover_moe.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "../../global_state/forward_context.hpp" +#include "../qwen3_next/qwen3_next_sparse_moe_block.hpp" + +namespace infinilm::models::clover_moe { + +class CloverSparseMoe : public qwen3_next::Qwen3NextSparseMoeBlock { +public: + CloverSparseMoe(runtime::Runtime &runtime, std::shared_ptr config, + size_t layer, const engine::distributed::RankInfo &rank, + backends::AttentionBackend backend, const layers::moe::EPConfig &ep); + runtime::Tensor forward(runtime::Runtime &runtime, const runtime::Tensor &hidden, + const global_state::ForwardContext &context) const; + +private: + std::tuple route(runtime::Runtime &runtime, + const runtime::Tensor &hidden, const global_state::ForwardContext &context) const; + runtime::Tensor router_weight_; + mutable runtime::Tensor cached_routed_; + bool coreset_; + double top_p_; + size_t top_k_; +}; + +} // namespace infinilm::models::clover_moe diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index 073bd4681..f35e38bfb 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -67,6 +67,8 @@ class InfinilmModel : public infinilm::weights::ModelWeights { /// Preserve logits for every packed position for speculative/MTP callers. bool sample_all_positions{false}; size_t logical_kv_upper_bound{0}; + bool diffusion_preview{false}; + bool diffusion_expert_budget{false}; }; struct Output { diff --git a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp index a74578fa9..aba19843d 100644 --- a/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp +++ b/csrc/models/qwen3_next/qwen3_next_gated_deltanet.hpp @@ -64,7 +64,7 @@ class Qwen3NextGatedDeltaNet : public infinilm::weights::ModelWeights { infinilm::runtime::Tensor forward(const infinilm::runtime::Tensor &hidden_states, const infinilm::global_state::ForwardContext &forward_context) const; -private: +protected: std::shared_ptr in_proj_qkv_; std::shared_ptr in_proj_z_; std::shared_ptr in_proj_a_; diff --git a/csrc/pybind11/bindings.cc b/csrc/pybind11/bindings.cc index 4a4dda6aa..2ef7d37c3 100644 --- a/csrc/pybind11/bindings.cc +++ b/csrc/pybind11/bindings.cc @@ -37,6 +37,10 @@ PYBIND11_MODULE(_infinilm, m) { || !checkpoint_float(input->dtype()) || !checkpoint_float(output->dtype())) { throw std::invalid_argument("checkpoint Cast requires matching CPU shapes and FP32/FP16/BF16 tensors"); } + if (input->numel() != 0 && input->data() == output->data() + && input->dtype() != output->dtype()) { + throw std::invalid_argument("checkpoint dtype conversion requires distinct storage"); + } const infini::ops::Tensor input_view{input->view()}; const infini::ops::Tensor output_view{output->view()}; infini::ops::Cast::Call(input_view, output_view); }, py::arg("input"), py::arg("output"), py::call_guard()); diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index 2fbcdeec1..5927ecfe4 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -233,6 +233,8 @@ inline void bind_infer_engine(py::module &m) { "random_values", "sampling_rows", "nsa_owner_snapshot", + "diffusion_preview", + "diffusion_expert_budget", }; for (auto &item : kwargs) { @@ -243,7 +245,11 @@ inline void bind_infer_engine(py::module &m) { "InferEngine.Input got an unexpected keyword argument '" + key + "'"); } - if (key == "temperature") { + if (key == "diffusion_expert_budget") { + input.diffusion_expert_budget = py::cast(item.second); + } else if (key == "diffusion_preview") { + input.diffusion_preview = py::cast(item.second); + } else if (key == "temperature") { input.temperature = py::cast(item.second); } else if (key == "top_p") { input.top_p = py::cast(item.second); @@ -303,6 +309,8 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) + .def_readwrite("diffusion_preview", &InferEngine::Input::diffusion_preview) + .def_readwrite("diffusion_expert_budget", &InferEngine::Input::diffusion_expert_budget) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p) diff --git a/docs/clover-hygon-ops.cmake b/docs/clover-hygon-ops.cmake new file mode 100644 index 000000000..9dd4023bf --- /dev/null +++ b/docs/clover-hygon-ops.cmake @@ -0,0 +1,4 @@ +# Public APIs needed to compile the complete InfiniLM source tree. +set(INFINI_OPS_DECLARATIONS + "add;argmax;cast;cat;causal_conv1d;chunk_gated_delta_rule;clover_block_conv1d;clover_coreset;clover_local_memory;clover_recurrent_gdn;convolution;copy;deepseek_moe;embedding;fill;flash_attn_varlen_func;flash_attn_with_kvcache;fused_add_rms_norm;fused_gated_delta_net_gating;gelu;gemm;i8_gemm;internal_softmax;layer_norm;mamba_selective_scan;moe_align_block_size;moe_fused_dense;moe_fused_gate;mrope;mul;pad;per_channel_quant_i8;per_tensor_dequant_i8;per_tensor_quant_i8;quick_gelu;recurrent_gated_delta_rule;relu;reshape_and_cache_flash;rms_norm;rotary_embedding;row_l2_normalize;separate_table_rope;sigmoid;silu;silu_and_mul;softmax;sum;take;top_k_top_p_sampling_from_logits;topk_sigmoid;topk_softmax;tril;triu;upsample_bicubic2d;upsample_bilinear2d;w4_marlin_gemm;w4_marlin_parameters;w4_marlin_repack" + CACHE STRING "Additional public operator declarations" FORCE) diff --git a/docs/clover-hygon.md b/docs/clover-hygon.md new file mode 100644 index 000000000..05ef81f80 --- /dev/null +++ b/docs/clover-hygon.md @@ -0,0 +1,169 @@ +# Clover block diffusion on Hygon + +`clover_moe` loads the complete Clover-35B-A3B-Preview checkpoint, including its +right-convolution and bidirectional GDN scalar shards. Model forward runs in +InfiniLM C++; InfiniOps supplies Hygon kernels. Triton compiles kernels on first +use and the native driver launches their binaries thereafter. Torch is used by +the Python decoder for confidence selection, not as a model-forward fallback. + +The model currently requires BF16 weights. The validated configuration is TP=2 on two Hygon BW 64GB devices, DTK +26.04, eager execution, static cache, and 32-token diffusion blocks. The chat +endpoint serializes requests and supports greedy, non-streaming text generation. +The measurements below cover concurrency one and short prompts; they are not a +measurement of maximum concurrent-serving capacity or 128K-context performance. + +## Build + +Build the matching InfiniRT, InfiniOps and InfiniCCL revisions first. The Hygon +InfiniRT driver API and InfiniOps Clover providers are required; the original +InfiniLM dependency pin does not supply them. The dependency PR links are +recorded in this change's PR description. + +```sh +export INFINI_ROOT=/path/to/installed-stack +export CUDA_HOME=/opt/dtk/cuda/cuda +export LD_LIBRARY_PATH="$INFINI_ROOT/lib:$CUDA_HOME/lib64:$LD_LIBRARY_PATH" +cmake -S third_party/InfiniCore/submodules/InfiniRT -B build/rt \ + -DCMAKE_INSTALL_PREFIX="$INFINI_ROOT" -DWITH_CPU=ON -DWITH_HYGON=ON +cmake --build build/rt -j 12 +cmake --install build/rt +cmake -S third_party/InfiniCore/submodules/InfiniOps -B build/ops \ + -C docs/clover-hygon-ops.cmake -DCMAKE_INSTALL_PREFIX="$INFINI_ROOT" \ + -DINFINI_RT_ROOT="$INFINI_ROOT" -DWITH_CPU=ON -DWITH_HYGON=ON \ + -DWITH_TRITON=ON -DHYGON_ARCH=gfx936 -DGENERATE_PYTHON_BINDINGS=ON \ + -DAUTO_DETECT_DEVICES=OFF -DAUTO_DETECT_BACKENDS=OFF +cmake --build build/ops -j 12 +cmake --install build/ops +cmake -S third_party/InfiniCore/submodules/InfiniCCL -B build/ccl \ + -DCMAKE_INSTALL_PREFIX="$INFINI_ROOT" -DWITH_HYGON=ON -DWITH_NCCL=ON \ + -DAUTO_DETECT_DEVICES=OFF -DBUILD_EXAMPLES=OFF +cmake --build build/ccl -j 12 +cmake --install build/ccl +export INFINILM_CXX11_ABI=1 +xmake f -y -m release --cxx11-abi=1 +pip install -e '.[clover]' --no-build-isolation +``` + +Match the C++ ABI to the installed DTK Torch build. `gfx936` is the tested +architecture, not a default for all Hygon products. Install the InfiniOps Python +package as `infini` as well: its JIT compiler and kernel sources must remain +importable. A direct CMake install puts that package at the install prefix; +placing a link named `infini` to this prefix in a Python search directory is +sufficient for a development build. A wheel installs this namespace normally. + +The declaration preset lets the complete InfiniLM source compile when other +model families reference operators without Hygon providers. Such operators +remain unsupported; generating their APIs does not install substitute kernels. + +## Generate and serve + +```python +from infinilm.diffusion.clover import CloverGenerator + +generator = CloverGenerator('/path/to/Clover-35B-A3B-Preview', tp=2, max_context=2048) +result = generator.generate('What is the capital of France?', max_new_tokens=64) +print(result.text) +``` + +```sh +HIP_VISIBLE_DEVICES=0,1 python -m infinilm.server.clover_server \ + --model /path/to/Clover-35B-A3B-Preview --tp 2 --max-context 2048 \ + --policy frontier --expert-budget --frontier-floor 0.7 --port 8000 +``` + +Send ordinary `/v1/chat/completions` requests with `temperature: 0` and +`stream: false`. Optional request fields are `diffusion_policy` (`baseline` or +`frontier`), `expert_budget`, and `frontier_confidence_floor`. Neutral repetition, +presence and frequency penalties and thinking disabled define the tested lane. +Nonzero penalties, stochastic sampling and streaming receive explicit errors. + +## Methods and state contract + +Preview reads the committed convolution/GDN state. Full attention can read the +causal prefix and all positions within the current block. Preview KV writes are +confined to the uncommitted suffix. Every block is fully resolved before causal +replay updates persistent state. Mask tokens are never emitted or committed. + +`baseline` uses the checkpoint's 0.95 low-confidence rule with top-one progress +when no masked position clears the threshold. `frontier` adapts the target +stride from Ringssss/dLLM-serve V3. It retains all positions above 0.95, adds +safe positions toward the stride, and guarantees progress. Its confidence +floor is explicit because lowering it changes decoding and can harm output. + +`expert_budget` adapts the supplied dllm-h100 study's top-4 expert budget and +TV6 decoded-row reuse with refresh interval five. Clover's own block coreset +continues to define the allowed experts. A token is eligible for reuse only +when it was unmasked in both consecutive previews. Inactive routes are omitted +from expert alignment/GEMMs and their previously combined outputs are restored. +Shared experts and causal replay still compute normally. State is reset at each +new block/request. These are optional approximate preview optimizations. + +This port does not implement the H100 code's DP/EP sequence-parallel communication +pipeline, or dLLM-serve's complete admission/compaction scheduler. The serial +endpoint exposes the mechanisms validated here. It also does not implement +Clover's separate self-speculative AR-verification serving mode. + +Source provenance: + +- dLLM-serve: `1da50a96c9ce6b888d245d378f2bac0425feba8c`, + `cwsrpt_v2/algorithm/cw_srpt_v3.py`, Apache-2.0; retained license/notice under + `python/infinilm/licenses/dLLM-serve`. +- Supplied dllm-h100 archive: top-4 routing and refresh/live-row logic in + `codex_coding/src/{test_m_skip_sweep.py,tv6_patch.py}`; reimplemented for the + Clover softmax router and InfiniLM state ownership. +- Preview kernels are provided and attributed by the companion InfiniOps PR. + +## Validation and measurement + +The native baseline matched the local SGLang Clover port's output tokens on +four initial prompts. Operator comparisons cover preview convolution/GDN, +clean-state immutability, packed recurrent commits, expert selection, fused +MoE, liveness and cached-row merging. Broader generative outputs can differ +across kernel implementations even with identical sampling settings. + +The recorded HTTP campaign uses the same model, tokenizer, BF16 precision, +TP2 budget, prompts, per-request output caps and sampling controls. Each of five +lanes ran 32 prompts three times after a separate warmup, alternating lane +order. Throughput is visible completion tokens, retokenized identically, divided +by summed client request time. SGLang's diffusion usage counters can include +block padding, so they are not used as the numerator. Results and the exact +workload are recorded in the PR evidence; no profiling run contributes to timing. + +A portable client and the controlled 32-prompt workload are included: + +```sh +python examples/bench_clover.py --url http://127.0.0.1:8000 --model clover \ + --tokenizer /path/to/Clover-35B-A3B-Preview --output native.jsonl +python examples/bench_clover.py --url http://127.0.0.1:8001 --model /served/model/name \ + --tokenizer /path/to/Clover-35B-A3B-Preview --backend sglang --variants baseline --output sglang.jsonl +``` + +Repeat `-1` is warmup and must be excluded. Sum `visible_tokens` over successful +requests and divide by summed `elapsed_s` over all measured requests, retaining +failures. Keep model weights, TP, sampling, output caps and context capacity fixed. + +### Recorded short-prompt campaign + +| Variant | Visible output tok/s | Per DCU tok/s | +| --- | ---: | ---: | +| InfiniLM baseline | 28.57 | 14.29 | +| Frontier stride | 38.47 | 19.24 | +| Expert budget and live-row reuse | 28.70 | 14.35 | +| Joint, confidence floor 0.7 | 39.92 | 19.96 | +| Local SGLang Clover iterative port | 5.88 | 2.94 | + +The joint path is 1.397x the native baseline and 6.791x this local SGLang port. +Three joint repetitions measured 39.85–40.02 visible tok/s. The SGLang reference +is the local DTK 0.5.12 port of iterative Clover, with eager Triton attention/MoE; +it is not the official self-speculative implementation. The two lanes use +separate pairs of identical BW devices on the same host. Context capacity is +2048 for both. These numbers include HTTP/request latency and real EOS behavior. + +All requests completed. The controlled structural/answer checks passed 96/96 +for joint, 93/96 for expert-only and 90/96 for baseline/frontier/SGLang. This +checks the included fixture, not general quality equivalence. Generated text +can still contain imperfect wording. A lower floor of 0.5 produced invalid +Python and repetitions and was rejected. Keep outputs and failed candidates in +any further evaluation. + +![Rendered validation record](images/clover-hygon-validation.png) diff --git a/docs/images/clover-hygon-validation.png b/docs/images/clover-hygon-validation.png new file mode 100644 index 000000000..61c415b6a Binary files /dev/null and b/docs/images/clover-hygon-validation.png differ diff --git a/examples/bench_clover.py b/examples/bench_clover.py new file mode 100644 index 000000000..e77f91641 --- /dev/null +++ b/examples/bench_clover.py @@ -0,0 +1,103 @@ +"""Measure serial Clover chat requests and retain every response and failure.""" + +import argparse +import json +import time +import urllib.error +import urllib.request +from pathlib import Path + +from transformers import AutoTokenizer + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--url", required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--backend", choices=("infinilm", "sglang"), default="infinilm") + parser.add_argument("--tokenizer", required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--workload", + type=Path, + default=Path(__file__).parents[1] / "test/diffusion/workload.json", + ) + parser.add_argument("--variants", default="baseline,frontier,budget,joint") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--frontier-floor", type=float, default=0.7) + args = parser.parse_args() + variants = args.variants.split(",") + if not variants or set(variants) - {"baseline", "frontier", "budget", "joint"}: + parser.error("Unknown variant.") + if args.backend == "sglang" and variants != ["baseline"]: + parser.error("The SGLang comparison requires --variants baseline.") + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer, trust_remote_code=True) + workload = json.loads(args.workload.read_text()) + args.output.parent.mkdir(parents=True, exist_ok=True) + for repeat in range(-1, args.repeats): + for item in workload: + order = variants if (repeat + item["id"]) % 2 == 0 else variants[::-1] + for variant in order: + request = { + "model": args.model, + "messages": [{"role": "user", "content": item["prompt"]}], + "temperature": 0, + "max_tokens": item["max_tokens"], + "presence_penalty": 0, + "frequency_penalty": 0, + "repetition_penalty": 1, + "chat_template_kwargs": {"enable_thinking": False}, + } + if args.backend == "infinilm": + request.update( + diffusion_policy="frontier" + if variant in ("frontier", "joint") + else "baseline", + expert_budget=variant in ("budget", "joint"), + frontier_confidence_floor=args.frontier_floor, + ) + record = { + "repeat": repeat, + "prompt_id": item["id"], + "variant": variant, + "request": request, + "success": False, + } + start = time.perf_counter() + try: + http_request = urllib.request.Request( + args.url.rstrip("/") + "/v1/chat/completions", + data=json.dumps(request).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(http_request, timeout=240) as response: + result = json.load(response) + record["elapsed_s"] = time.perf_counter() - start + text = result["choices"][0]["message"]["content"] + record.update( + success=True, + response=result, + visible_tokens=len( + tokenizer(text, add_special_tokens=False)["input_ids"] + ), + ) + except (urllib.error.URLError, ValueError, KeyError) as error: + record.update( + elapsed_s=time.perf_counter() - start, + error=str(error), + visible_tokens=0, + ) + with args.output.open("a") as output: + output.write(json.dumps(record, ensure_ascii=False) + "\n") + print( + variant, + repeat, + item["id"], + record["success"], + round(record["elapsed_s"], 4), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 8c02f7fba..9dc7b3df7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ classifiers = [ [project.optional-dependencies] torch-compat = ["torch"] +clover = ["torch", "triton", "fastapi", "uvicorn"] [project.urls] Homepage = "https://github.com/InfiniTensor/InfiniLM" diff --git a/python/infinilm/diffusion/__init__.py b/python/infinilm/diffusion/__init__.py new file mode 100644 index 000000000..a7fac7008 --- /dev/null +++ b/python/infinilm/diffusion/__init__.py @@ -0,0 +1 @@ +"""Block diffusion inference with explicit preview and causal commit phases.""" diff --git a/python/infinilm/diffusion/clover.py b/python/infinilm/diffusion/clover.py new file mode 100644 index 000000000..0682554c7 --- /dev/null +++ b/python/infinilm/diffusion/clover.py @@ -0,0 +1,181 @@ +"""Native Clover block diffusion on the InfiniLM C++ model and InfiniOps.""" + +import time +from dataclasses import dataclass + +from .frontier import FrontierPolicy + + +@dataclass +class CloverResult: + text: str + token_ids: list[int] + iterations: list[int] + elapsed_s: float + + @property + def output_tokens_per_second(self): + return len(self.token_ids) / self.elapsed_s + + +class CloverGenerator: + """Run greedy diffusion with clean prefix state and causal block replay.""" + + def __init__(self, model_path, *, tp=2, max_context=4096): + import torch + from transformers import AutoTokenizer + + from infinilm import _core as core + from infinilm.cache import StaticKVCacheConfig + from infinilm.distributed import DistConfig + from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + + self.torch = torch + self.core = core + self.cache_config = StaticKVCacheConfig(1, max_context) + self.engine = InferEngine( + model_path, + device=core.device("hygon"), + distributed_config=DistConfig(tp), + cache_config=self.cache_config, + ) + if self.engine.model_type != "clover_moe": + raise ValueError("CloverGenerator requires a Clover checkpoint.") + load_model_state_dict_by_file(self.engine, model_path, self.engine.dtype) + self.tokenizer = AutoTokenizer.from_pretrained( + model_path, trust_remote_code=True + ) + self.mask_id = self.engine.hf_config["mask_token_id"] + self.block_size = self.engine.hf_config["block_size"] + self.max_context = max_context + + def _tensor(self, values, *, wide=False): + import numpy as np + + return self.core.from_numpy( + np.asarray(values, dtype=np.int64 if wide else np.int32) + ) + + def _forward(self, ids, past, *, preview=False, expert_budget=False): + length = len(ids) + return self.engine.forward_raw( + self._tensor([ids], wide=True), + position_ids=self._tensor([list(range(past, past + length))], wide=True), + past_kv_lengths=self._tensor([past]), + total_kv_lengths=self._tensor([past + length]), + input_offsets=self._tensor([0, length]), + cu_seqlens=self._tensor([0, past + length]), + block_tables=self._tensor([[0]]), + mamba_init_state_indices=self._tensor([0]), + mamba_final_state_indices=self._tensor([0]), + diffusion_preview=preview, + diffusion_expert_budget=expert_budget, + sample_all_positions=preview, + temperature=0, + top_k=1, + ) + + def _logits(self, value): + torch, core = self.torch, self.core + out = torch.empty(value.shape, device="cuda:0", dtype=torch.bfloat16) + target = core.from_blob( + out.data_ptr(), + list(value.shape), + dtype=core.bfloat16, + device=core.device("hygon", 0), + owner=out, + ) + target.copy_(value) + core.set_device(core.device("hygon", 0)) + core.sync_device() + return out.view(self.block_size, -1).float() + + def generate( + self, + prompt, + *, + max_new_tokens=128, + policy="baseline", + threshold=0.95, + ignore_eos=False, + expert_budget=False, + frontier_confidence_floor=0.9, + ): + if policy not in ("baseline", "frontier"): + raise ValueError("Policy must be baseline or frontier.") + if max_new_tokens < 0 or not 0 <= threshold <= 1: + raise ValueError("Invalid output limit or confidence threshold.") + torch = self.torch + prompt_ids = self.tokenizer.apply_chat_template( + ( + [{"role": "user", "content": prompt}] + if isinstance(prompt, str) + else prompt + ), + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ) + if hasattr(prompt_ids, "keys"): + prompt_ids = prompt_ids["input_ids"] + blocks = (max_new_tokens + self.block_size - 1) // self.block_size + if len(prompt_ids) + blocks * self.block_size > self.max_context: + raise ValueError( + "Prompt and complete diffusion blocks exceed the cache capacity." + ) + self.engine.reset_cache(self.cache_config) + start = time.perf_counter() + if max_new_tokens == 0: + return CloverResult("", [], [], time.perf_counter() - start) + self._forward(prompt_ids, 0) + output, iterations = [], [] + frontier = FrontierPolicy(confidence_floor=frontier_confidence_floor) + eos = set(self.engine.eos_token_id) + for _ in range(blocks): + ids = torch.full( + (self.block_size,), self.mask_id, device="cuda:0", dtype=torch.int64 + ) + for iteration in range(self.block_size): + result = self._forward( + ids.cpu().tolist(), + len(prompt_ids) + len(output), + preview=True, + expert_budget=expert_budget, + ) + logits = self._logits(result["logits"]) + logits[:, self.mask_id] = -float("inf") + confidence, proposed = logits.softmax(-1).max(-1) + masked = ids == self.mask_id + if policy == "frontier": + selected = frontier.select( + confidence[None], masked[None], iteration + )[0] + else: + selected = masked & (confidence >= threshold) + if not bool(selected.any()): + selected[confidence.masked_fill(~masked, -1).argmax()] = True + ids = torch.where(selected, proposed, ids) + if not bool((ids == self.mask_id).any()): + break + resolved = ids.cpu().tolist() + if self.mask_id in resolved: + raise RuntimeError("An unresolved diffusion block cannot be committed.") + self._forward(resolved, len(prompt_ids) + len(output)) + iterations.append(iteration + 1) + remaining = max_new_tokens - len(output) + for token in resolved[:remaining]: + if not ignore_eos and token in eos: + return CloverResult( + self.tokenizer.decode(output, skip_special_tokens=True), + output, + iterations, + time.perf_counter() - start, + ) + output.append(token) + return CloverResult( + self.tokenizer.decode(output, skip_special_tokens=True), + output, + iterations, + time.perf_counter() - start, + ) diff --git a/python/infinilm/diffusion/frontier.py b/python/infinilm/diffusion/frontier.py new file mode 100644 index 000000000..d1af1a938 --- /dev/null +++ b/python/infinilm/diffusion/frontier.py @@ -0,0 +1,43 @@ +"""Frontier-guided denoising adapted from Ringssss/dLLM-serve (Apache-2.0).""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class FrontierPolicy: + """Choose confident positions with a bounded, adaptive target stride.""" + + target_iterations: int = 8 + min_unmask: int = 2 + max_unmask: int = 16 + confidence_floor: float = 0.9 + confident_threshold: float = 0.95 + + def __post_init__(self): + if not 1 <= self.min_unmask <= self.max_unmask <= 32: + raise ValueError("Unmask bounds must satisfy 1 <= min <= max <= 32.") + if self.target_iterations < 1 or not ( + 0 <= self.confidence_floor <= self.confident_threshold <= 1 + ): + raise ValueError("Invalid target iterations or confidence floor.") + + def select(self, confidence, masked, iteration): + import torch + + remaining = masked.sum(-1) + rounds = max(1, self.target_iterations - iteration) + stride = torch.div(remaining + rounds - 1, rounds, rounding_mode="floor") + stride = torch.minimum( + stride.clamp(self.min_unmask, self.max_unmask), remaining + ) + scores = confidence.masked_fill(~masked, -float("inf")) + order = scores.argsort(dim=-1, descending=True, stable=True) + ranks = torch.arange(masked.shape[-1], device=masked.device) + take = ranks < stride[..., None] + transfer = torch.zeros_like(masked).scatter(-1, order, take) + transfer &= masked & (confidence >= self.confidence_floor) + transfer |= masked & (confidence >= self.confident_threshold) + stalled = masked.any(-1) & ~transfer.any(-1) + best = scores.argmax(-1, keepdim=True) + transfer.scatter_(-1, best, transfer.gather(-1, best) | stalled[..., None]) + return transfer diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 67cc645ca..05a8c29f8 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -338,6 +338,8 @@ def _build_input( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=False, + diffusion_preview=False, + diffusion_expert_budget=False, temperature=None, top_k=None, top_p=None, @@ -399,6 +401,8 @@ def convert_tensor_list(tensor_list_): visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + diffusion_preview=diffusion_preview, + diffusion_expert_budget=diffusion_expert_budget, temperature=temperature, top_k=top_k, top_p=top_p, @@ -534,6 +538,10 @@ def forward_raw( visual_token_ranges=None, target_hidden_states=None, sample_all_positions=True, + diffusion_preview=False, + diffusion_expert_budget=False, + mamba_init_state_indices=None, + mamba_final_state_indices=None, temperature=None, top_k=None, top_p=None, @@ -557,6 +565,10 @@ def forward_raw( visual_token_ranges=visual_token_ranges, target_hidden_states=target_hidden_states, sample_all_positions=sample_all_positions, + diffusion_preview=diffusion_preview, + diffusion_expert_budget=diffusion_expert_budget, + mamba_init_state_indices=mamba_init_state_indices, + mamba_final_state_indices=mamba_final_state_indices, temperature=temperature, top_k=top_k, top_p=top_p, diff --git a/python/infinilm/licenses/dLLM-serve/LICENSE b/python/infinilm/licenses/dLLM-serve/LICENSE new file mode 100644 index 000000000..261eeb9e9 --- /dev/null +++ b/python/infinilm/licenses/dLLM-serve/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/python/infinilm/licenses/dLLM-serve/NOTICE b/python/infinilm/licenses/dLLM-serve/NOTICE new file mode 100644 index 000000000..8f60423ca --- /dev/null +++ b/python/infinilm/licenses/dLLM-serve/NOTICE @@ -0,0 +1,4 @@ +Frontier-guided stride control is adapted from Ringssss/dLLM-serve, commit +1da50a96c9ce6b888d245d378f2bac0425feba8c, cwsrpt_v2/algorithm/cw_srpt_v3.py. +The adaptation preserves confident-token progress and uses a configurable +confidence floor. Original source is distributed under Apache-2.0. diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 114aae846..b1916c66d 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -328,6 +328,7 @@ def load_model_state_dict_by_file( _remap_qwen3_next, _remap_qwen3_5, _remap_qwen3_5_moe, + _remap_clover_moe, _remap_kimi_k3, ) or ( @@ -1328,6 +1329,12 @@ def _remap_kimi_k3(state_dict, config): return state_dict +def _remap_clover_moe(state_dict, config, *, defer_tie=False): + """Load Clover's Qwen weights while retaining every diffusion parameter.""" + wrapped = dict(config, text_config=config) + return _remap_qwen3_5_moe(state_dict, wrapped, defer_tie=defer_tie) + + _WEIGHT_REMAPPER = { "glm4": _remap_glm4, "chatglm": _remap_chatglm, @@ -1337,6 +1344,7 @@ def _remap_kimi_k3(state_dict, config): "videonsa": _remap_videonsa, "qwen3_5": _remap_qwen3_5, "qwen3_5_moe": _remap_qwen3_5_moe, + "clover_moe": _remap_clover_moe, "qwen3_next": _remap_qwen3_next, "kimi_k3": _remap_kimi_k3, } diff --git a/python/infinilm/server/clover_server.py b/python/infinilm/server/clover_server.py new file mode 100644 index 000000000..338e1bf25 --- /dev/null +++ b/python/infinilm/server/clover_server.py @@ -0,0 +1,146 @@ +"""OpenAI-compatible serial serving for the native Clover diffusion engine.""" + +import argparse +import asyncio +import time +import uuid + +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel, Field + +from infinilm.diffusion.clover import CloverGenerator + + +class ChatRequest(BaseModel): + model: str = "clover" + messages: list[dict] + max_tokens: int = Field(default=128, ge=0) + temperature: float = 0 + stream: bool = False + diffusion_policy: str | None = None + expert_budget: bool | None = None + frontier_confidence_floor: float | None = Field(default=None, ge=0, le=1) + presence_penalty: float = 0 + frequency_penalty: float = 0 + repetition_penalty: float = 1 + chat_template_kwargs: dict = Field(default_factory=dict) + + +def create_app( + generator, *, policy="baseline", expert_budget=False, frontier_floor=0.9 +): + """Serialize requests because each worker owns one recurrent state slot.""" + lock = asyncio.Lock() + + app = FastAPI() + + @app.get("/health") + def health(): + return {"status": "ok"} + + @app.get("/v1/models") + def models(): + return {"object": "list", "data": [{"id": "clover", "object": "model"}]} + + @app.post("/v1/chat/completions") + async def chat(request: ChatRequest): + if ( + request.presence_penalty != 0 + or request.frequency_penalty != 0 + or request.repetition_penalty != 1 + or request.chat_template_kwargs.get("enable_thinking", False) + ): + raise HTTPException( + 400, "This endpoint requires neutral penalties and thinking disabled." + ) + if request.stream or request.temperature != 0: + raise HTTPException( + 400, "Clover currently supports non-streaming greedy requests." + ) + if any( + message.get("role") not in ("system", "user", "assistant") + or not isinstance(message.get("content"), str) + for message in request.messages + ): + raise HTTPException( + 400, "Messages require a supported role and text content." + ) + if not request.messages: + raise HTTPException(400, "At least one message is required.") + async with lock: + try: + result = await asyncio.to_thread( + generator.generate, + request.messages, + max_new_tokens=request.max_tokens, + policy=request.diffusion_policy or policy, + frontier_confidence_floor=frontier_floor + if request.frontier_confidence_floor is None + else request.frontier_confidence_floor, + expert_budget=expert_budget + if request.expert_budget is None + else request.expert_budget, + ) + except ValueError as exc: + raise HTTPException(400, str(exc)) from exc + prompt = generator.tokenizer.apply_chat_template( + request.messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + prompt_count = len( + generator.tokenizer(prompt, add_special_tokens=False)["input_ids"] + ) + output_count = len(result.token_ids) + return { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": int(time.time()), + "model": request.model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": result.text}, + "finish_reason": "length" + if output_count >= request.max_tokens + else "stop", + } + ], + "usage": { + "prompt_tokens": prompt_count, + "completion_tokens": output_count, + "total_tokens": prompt_count + output_count, + }, + } + + return app + + +def main(): + import uvicorn + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True) + parser.add_argument("--tp", type=int, default=2) + parser.add_argument("--max-context", type=int, default=4096) + parser.add_argument( + "--policy", choices=("baseline", "frontier"), default="baseline" + ) + parser.add_argument("--expert-budget", action="store_true") + parser.add_argument("--frontier-floor", type=float, default=0.9) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + args = parser.parse_args() + generator = CloverGenerator(args.model, tp=args.tp, max_context=args.max_context) + app = create_app( + generator, + policy=args.policy, + expert_budget=args.expert_budget, + frontier_floor=args.frontier_floor, + ) + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/test/diffusion/test_clover_loading.py b/test/diffusion/test_clover_loading.py new file mode 100644 index 000000000..e7d5d2772 --- /dev/null +++ b/test/diffusion/test_clover_loading.py @@ -0,0 +1,38 @@ +"""Protect the text-only Clover checkpoint layout and learned preview weights.""" + +import numpy as np +from infinilm import _core as core +from infinilm.modeling_utils import _remap_clover_moe + + +def test_clover_remap_preserves_preview_weights_and_splits_qkv(): + config = { + "linear_key_head_dim": 2, + "linear_num_key_heads": 1, + "num_experts": 2, + "moe_intermediate_size": 2, + } + prefix = "model.layers.0.linear_attn." + norm = core.from_list([0.25, -0.5], dtype=core.float32) + scalar = core.from_list([5.0, 6.0], dtype=core.float32) + packed = core.from_numpy(np.arange(24, dtype=np.float32).reshape(6, 4)) + output = _remap_clover_moe( + { + "model.norm.weight": norm, + prefix + "norm.weight": norm, + prefix + "bd_softmax_history_bias": scalar, + prefix + "in_proj_qkv.weight": packed, + }, + config, + ) + np.testing.assert_array_equal(output["model.norm.weight"].to_numpy(), [1.25, 0.5]) + np.testing.assert_array_equal( + output[prefix + "norm.weight"].to_numpy(), [0.25, -0.5] + ) + assert output[prefix + "bd_softmax_history_bias"] is scalar + assert prefix + "in_proj_qkv.weight" not in output + for part, begin in (("q", 0), ("k", 2), ("v", 4)): + np.testing.assert_array_equal( + output[prefix + f"in_proj_{part}.weight"].to_numpy(), + packed.to_numpy()[begin : begin + 2], + ) diff --git a/test/diffusion/test_frontier.py b/test/diffusion/test_frontier.py new file mode 100644 index 000000000..665d17f4c --- /dev/null +++ b/test/diffusion/test_frontier.py @@ -0,0 +1,35 @@ +"""Exercise progress, resolved-token preservation and high-confidence bypass.""" + +import importlib.util +import sys +from pathlib import Path + +import torch + +spec = importlib.util.spec_from_file_location( + "clover_frontier", + Path(__file__).parents[2] / "python/infinilm/diffusion/frontier.py", +) +module = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = module +spec.loader.exec_module(module) + + +def test_confident_tokens_are_not_delayed_by_stride(): + policy = module.FrontierPolicy() + confidence = torch.ones(2, 32) + masked = torch.ones(2, 32, dtype=torch.bool) + masked[1, :7] = False + assert torch.equal(policy.select(confidence, masked, 0), masked) + + +def test_stalled_rows_progress_without_touching_finished_rows(): + policy = module.FrontierPolicy() + confidence = torch.full((3, 32), 0.1) + masked = torch.zeros(3, 32, dtype=torch.bool) + masked[0, 5:] = True + masked[1, 17] = True + selected = policy.select(confidence, masked, 3) + assert selected.sum(-1).tolist() == [1, 1, 0] + assert not (selected & ~masked).any() + assert selected[0, 5] and selected[1, 17] diff --git a/test/diffusion/workload.json b/test/diffusion/workload.json new file mode 100644 index 000000000..883b3870c --- /dev/null +++ b/test/diffusion/workload.json @@ -0,0 +1,194 @@ +[ + { + "id": 0, + "prompt": "What is 2 + 2? Reply with only the number.", + "max_tokens": 32, + "check": "4" + }, + { + "id": 1, + "prompt": "What is the capital of France? Answer in one sentence.", + "max_tokens": 64, + "check": "Paris" + }, + { + "id": 2, + "prompt": "What is the capital of Germany? Answer briefly.", + "max_tokens": 64, + "check": "Berlin" + }, + { + "id": 3, + "prompt": "Sort these numbers in ascending order: 3, 1, 2. Reply only with a JSON array.", + "max_tokens": 64, + "check": "[1, 2, 3]" + }, + { + "id": 4, + "prompt": "Return only this JSON object: {\"ok\": true, \"count\": 3}", + "max_tokens": 64, + "check": "json" + }, + { + "id": 5, + "prompt": "Translate \"good morning\" into Chinese. Answer briefly.", + "max_tokens": 64, + "check": "早" + }, + { + "id": 6, + "prompt": "Reverse the string abc. Output only the reversed string.", + "max_tokens": 32, + "check": "cba" + }, + { + "id": 7, + "prompt": "What is 6 multiplied by 7? Answer with only the number.", + "max_tokens": 32, + "check": "42" + }, + { + "id": 8, + "prompt": "Write a Python function add(a, b) that returns their sum. Return only the function.", + "max_tokens": 128, + "check": "code" + }, + { + "id": 9, + "prompt": "Write a Python function square(x) that returns x squared. Return only the function.", + "max_tokens": 128, + "check": "code" + }, + { + "id": 10, + "prompt": "用一句话解释什么是机器学习。", + "max_tokens": 128, + "check": null + }, + { + "id": 11, + "prompt": "Explain photosynthesis in two short sentences.", + "max_tokens": 128, + "check": null + }, + { + "id": 12, + "prompt": "List three primary colors of light, with no explanation.", + "max_tokens": 64, + "check": "red" + }, + { + "id": 13, + "prompt": "Give two benefits of unit testing in software.", + "max_tokens": 128, + "check": null + }, + { + "id": 14, + "prompt": "用两句话解释为什么需要备份数据。", + "max_tokens": 128, + "check": null + }, + { + "id": 15, + "prompt": "Write a polite one-sentence message declining a meeting invitation.", + "max_tokens": 128, + "check": null + }, + { + "id": 16, + "prompt": "Summarize this in one sentence: Bees transfer pollen between flowers. This helps many plants produce fruit and seeds.", + "max_tokens": 128, + "check": null + }, + { + "id": 17, + "prompt": "What is the opposite of hot? Reply with one word.", + "max_tokens": 32, + "check": "cold" + }, + { + "id": 18, + "prompt": "Return a JSON object with name Ada and age 30. No explanation.", + "max_tokens": 64, + "check": "json" + }, + { + "id": 19, + "prompt": "Complete the sequence: 2, 4, 6, 8, __. Reply with only the next number.", + "max_tokens": 32, + "check": "10" + }, + { + "id": 20, + "prompt": "Write a Python function is_even(n) that returns whether n is even. Return only code.", + "max_tokens": 128, + "check": "code" + }, + { + "id": 21, + "prompt": "Explain the difference between a list and a tuple in Python in two sentences.", + "max_tokens": 128, + "check": null + }, + { + "id": 22, + "prompt": "请用一句话解释什么是开源软件。", + "max_tokens": 128, + "check": null + }, + { + "id": 23, + "prompt": "计算：一个盒子有12个苹果，拿走5个，再放入8个，还剩多少个？", + "max_tokens": 64, + "check": "15" + }, + { + "id": 24, + "prompt": "What is the largest planet in our solar system? Answer briefly.", + "max_tokens": 64, + "check": "Jupiter" + }, + { + "id": 25, + "prompt": "Convert 100 centimeters to meters. Answer briefly.", + "max_tokens": 32, + "check": "1" + }, + { + "id": 26, + "prompt": "Write a short definition of a database index.", + "max_tokens": 128, + "check": null + }, + { + "id": 27, + "prompt": "Translate \"Thank you for your help\" into Chinese.", + "max_tokens": 64, + "check": "帮助" + }, + { + "id": 28, + "prompt": "What is 10 minus 3? Reply with only the number.", + "max_tokens": 32, + "check": "7" + }, + { + "id": 29, + "prompt": "Name the chemical formula of water. Reply with only the formula.", + "max_tokens": 32, + "check": "H2O" + }, + { + "id": 30, + "prompt": "Write a Python function that returns the nth Fibonacci number.", + "max_tokens": 192, + "check": "code" + }, + { + "id": 31, + "prompt": "Describe how to make tea in three short steps.", + "max_tokens": 128, + "check": null + } +] diff --git a/third_party/InfiniCore b/third_party/InfiniCore index 2b4f2f205..6446bc6b9 160000 --- a/third_party/InfiniCore +++ b/third_party/InfiniCore @@ -1 +1 @@ -Subproject commit 2b4f2f20584a69b074f62a455e3de7368e3446c2 +Subproject commit 6446bc6b9e341038ba13da362c7aecd5fdbeebee