From fe0ca456f7cf662f2712cae5c6d5fc92dcaed0b8 Mon Sep 17 00:00:00 2001 From: yangyeyuan <2776079516@qq.com> Date: Sun, 20 Sep 2026 15:25:17 +0800 Subject: [PATCH 1/4] feat: add FP8 blockwise quantization for weights and KV cache Weight path: - add the fp8_blockwise quantization scheme and wire it into Linear layer construction and FP8 checkpoint loading (weight_scale_inv) - decode GEMM routing: fused fp8_blockwise_gemm kernel for small M, tensor-core mma path up to M=32, dequantize+GEMM fallback otherwise - opt-in Marlin fused GEMM path, with the rounding carry in F16 scale conversion KV cache path: - FP8(E4M3) KV cache quantization for the paged attention backend with per-token-per-kv-head dynamic scales - expose kv_cache_dtype=fp8 through EngineConfig, LLM and the server entry points; reject FP8 KV cache on non-paged attention backends Requires the InfiniCore fp8_blockwise_gemm / fp8_blockwise_dequantize operators and FP8 paged attention support. --- csrc/cache/kv_cache.cpp | 17 ++ csrc/cache/kv_cache.hpp | 11 + csrc/config/quant_config.cpp | 2 + csrc/config/quant_config.hpp | 4 + csrc/engine/infer_engine.cpp | 8 +- csrc/global_state/forward_context.hpp | 7 + csrc/layers/attention/attention.cpp | 4 + csrc/layers/attention/backends/flash_attn.cpp | 4 +- csrc/layers/attention/backends/paged_attn.cpp | 38 ++- csrc/layers/attention/backends/paged_attn.hpp | 5 +- csrc/layers/quantization/fp8_blockwise.cpp | 261 ++++++++++++++++++ csrc/layers/quantization/fp8_blockwise.hpp | 49 ++++ csrc/layers/quantization/gptq_marlin.cpp | 2 +- csrc/layers/quantization/gptq_marlin.hpp | 8 +- csrc/layers/quantization/marlin_utils.hpp | 3 + csrc/layers/quantization/quantization.hpp | 1 + .../quantization/quantization_scheme.hpp | 2 + .../deepseek_v2/deepseek_v2_mla_attention.cpp | 10 +- csrc/models/infinilm_model.cpp | 19 ++ csrc/models/videonsa/videonsa_attention.cpp | 2 +- python/infinilm/base_config.py | 2 +- python/infinilm/config/engine_config.py | 3 + python/infinilm/llm/llm.py | 8 + .../infinilm/llm/model_runner/model_runner.py | 1 + python/infinilm/modeling_utils.py | 74 ++++- python/infinilm/server/inference_server.py | 6 + python/infinilm/server/pipeline_worker.py | 1 + 27 files changed, 532 insertions(+), 20 deletions(-) create mode 100644 csrc/layers/quantization/fp8_blockwise.cpp create mode 100644 csrc/layers/quantization/fp8_blockwise.hpp diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp index 2a7780090..f88cab4e8 100644 --- a/csrc/cache/kv_cache.cpp +++ b/csrc/cache/kv_cache.cpp @@ -146,6 +146,23 @@ infinicore::Tensor create_layer_kv_cache( return kv_cache; } + +std::pair create_layer_kv_scales( + const infinicore::Size num_kv_heads, + const PagedKVCacheConfig &config) { + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + + // Mirror the KV-head sharding of create_layer_kv_cache (num_k_heads == num_v_heads there). + bool is_kv_replica = (num_kv_heads < rank_info.tp_size && rank_info.tp_size % num_kv_heads == 0); + size_t num_rank_kv_heads = is_kv_replica ? 1 : (num_kv_heads / rank_info.tp_size); + + // [num_blocks, num_rank_kv_heads, block_size], one F32 scale per token per kv head. + const infinicore::Shape scale_shape = {config.num_blocks(), num_rank_kv_heads, config.block_size()}; + infinicore::Tensor k_scale = infinicore::Tensor::zeros(scale_shape, infinicore::DataType::F32, rank_info.device); + infinicore::Tensor v_scale = infinicore::Tensor::zeros(scale_shape, infinicore::DataType::F32, rank_info.device); + + return {std::move(k_scale), std::move(v_scale)}; +} }; // namespace PagedKVCache } // namespace infinilm::cache diff --git a/csrc/cache/kv_cache.hpp b/csrc/cache/kv_cache.hpp index 760fb6685..66bf876df 100644 --- a/csrc/cache/kv_cache.hpp +++ b/csrc/cache/kv_cache.hpp @@ -5,6 +5,7 @@ #include #include +#include namespace infinilm::cache { class StaticKVCacheConfig final : public CacheConfig { @@ -62,6 +63,16 @@ infinicore::Tensor create_layer_kv_cache( infinicore::DataType dtype, const PagedKVCacheConfig &config); +// FP8(E4M3) KV cache scales for one layer: returns `{k_scale, v_scale}`, each F32 with +// shape `[num_blocks, num_rank_kv_heads, block_size]` (one scale per token per kv head; +// the head_dim dimension is reduced by the paged_caching kernel on write). The scale +// layout is logical (block, kv_head, token-in-block) and does not follow the cache +// tensor's element order, so it is identical for the paged (BHSD) and FLASH_ATTN (BSHD) +// cache layouts. +std::pair create_layer_kv_scales( + infinicore::Size num_kv_heads, + const PagedKVCacheConfig &config); + } // namespace PagedKVCache } // namespace infinilm::cache diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index e58966d89..6de8e8048 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -22,6 +22,8 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); } else if (quant_method == "quark") { return std::make_shared(quantization_config); + } else if (quant_method == "fp8") { + return std::make_shared(quantization_config); } else { return std::make_shared(quantization_config); } diff --git a/csrc/config/quant_config.hpp b/csrc/config/quant_config.hpp index fb0b8abf3..2dd203952 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -32,6 +32,10 @@ class QuantConfig { this->kv_quant_scheme = infinilm::quantization::KVQuantAlgo::INT8; break; } + case infinicore::DataType::F8: { + this->kv_quant_scheme = infinilm::quantization::KVQuantAlgo::FP8; + break; + } default: { spdlog::warn("Unsupported kv_cache_dtype: '{}', fallback to NONE", infinicore::toString(kv_cache_dtype)); this->kv_quant_scheme = infinilm::quantization::KVQuantAlgo::NONE; diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 422b5df73..537d51035 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -81,8 +81,14 @@ InferEngine::InferEngine( distributed_config.moe_ep_size, pre_transpose); - // Only support offline int8 kv cache quantization in this version + // KV cache quantization: INT8 is the offline static per-tensor path; FP8(E4M3) uses + // per-token-per-kv-head dynamic scales and currently requires the paged attention + // backend (STATIC_ATTN and FLASH_ATTN paths do not implement it). if (kv_cache_dtype.has_value()) { + if (kv_cache_dtype.value() == infinicore::DataType::F8 + && attention_backend != backends::AttentionBackend::PAGED_ATTN) { + throw std::invalid_argument("InferEngine: FP8 KV cache (kv_cache_dtype=fp8) requires the paged attention backend"); + } this->model_config_->set_kv_quant_scheme(kv_cache_dtype.value()); } // Create one RankWorker per rank diff --git a/csrc/global_state/forward_context.hpp b/csrc/global_state/forward_context.hpp index 6d02de4c5..1426610a1 100644 --- a/csrc/global_state/forward_context.hpp +++ b/csrc/global_state/forward_context.hpp @@ -2,6 +2,8 @@ #include "../models/infinilm_model.hpp" +#include + namespace infinilm::global_state { struct AttentionMetadata { @@ -68,6 +70,11 @@ struct ForwardContext { MambaMetadata mamba_metadata; MultiModalMetadata mm_metadata; std::vector kv_cache_vec; + // Per-layer FP8(E4M3) KV cache scales `{k_scale, v_scale}`, parallel to `kv_cache_vec`. + // Each is F32 with shape `[num_blocks, num_kv_heads, block_size]` (per-token-per-kv-head + // dynamic scale, filled by the paged_caching kernel on write). Empty when the KV cache + // dtype is not F8. + std::vector> kv_scale_vec; std::vector conv_state_vec; std::vector ssm_state_vec; }; diff --git a/csrc/layers/attention/attention.cpp b/csrc/layers/attention/attention.cpp index 16506ef02..72158009f 100644 --- a/csrc/layers/attention/attention.cpp +++ b/csrc/layers/attention/attention.cpp @@ -150,6 +150,10 @@ void init_kv_cache_quant_params(std::function FlashAttentionImpl::do_kv_cac v_cache_layer->permute({0, 2, 1, 3}), key, value, - slot_mapping); + slot_mapping, + std::nullopt, // no FP8 scales: FP8 KV cache is not supported on the FLASH_ATTN backend + std::nullopt); return {k_cache_layer, v_cache_layer}; } diff --git a/csrc/layers/attention/backends/paged_attn.cpp b/csrc/layers/attention/backends/paged_attn.cpp index f39937ead..99660bfd3 100644 --- a/csrc/layers/attention/backends/paged_attn.cpp +++ b/csrc/layers/attention/backends/paged_attn.cpp @@ -30,8 +30,28 @@ infinicore::Tensor PagedAttentionImpl::forward(const AttentionLayer &layer, ASSERT(block_tables.has_value()); ASSERT(slot_mapping.has_value()); + // FP8(E4M3) KV cache: fetch this layer's per-token-per-kv-head scales (allocated with + // the cache in ForwardContext::kv_scale_vec). Non-FP8 caches pass nullopt, keeping the + // operator behavior bitwise unchanged. + std::optional k_scale = std::nullopt; + std::optional v_scale = std::nullopt; + if (kv_cache->dtype() == infinicore::DataType::F8) { + const auto &kv_scale_vec = infinilm::global_state::get_forward_context().kv_scale_vec; + if (layer_idx_ >= kv_scale_vec.size() + || kv_scale_vec[layer_idx_].first.empty() + || kv_scale_vec[layer_idx_].second.empty()) { + throw std::runtime_error( + "infinilm::layers::attention::backends::PagedAttentionImpl: FP8 KV cache requires per-layer " + "k_scale/v_scale, but none were allocated for layer " + + std::to_string(layer_idx_) + + ". FP8 KV cache is currently supported only by the default paged KV cache allocation path."); + } + k_scale = kv_scale_vec[layer_idx_].first; + v_scale = kv_scale_vec[layer_idx_].second; + } + // 1. update paged kv cache - auto [k_total, v_total] = do_kv_cache_update(layer, key, value, kv_cache, slot_mapping.value()); + auto [k_total, v_total] = do_kv_cache_update(layer, key, value, kv_cache, slot_mapping.value(), k_scale, v_scale); size_t seq_len = query->shape()[0]; bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]); @@ -49,7 +69,9 @@ infinicore::Tensor PagedAttentionImpl::forward(const AttentionLayer &layer, total_sequence_lengths.value(), input_offsets.value(), std::nullopt, - scale_); + scale_, + k_scale, + v_scale); } else { infinicore::op::paged_attention_( attn_output, @@ -59,7 +81,9 @@ infinicore::Tensor PagedAttentionImpl::forward(const AttentionLayer &layer, block_tables.value(), total_sequence_lengths.value(), std::nullopt, - scale_); + scale_, + k_scale, + v_scale); } attn_output = attn_output->view({1, seq_len, num_heads_ * value_head_dim}); return attn_output; @@ -69,7 +93,9 @@ std::tuple PagedAttentionImpl::do_kv_cac const infinicore::Tensor key, const infinicore::Tensor value, infinicore::Tensor &kv_cache, - const infinicore::Tensor slot_mapping) const { + const infinicore::Tensor slot_mapping, + const std::optional &k_scale, + const std::optional &v_scale) const { auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); v_cache_layer = v_cache_layer->narrow({{3, 0, value->size(value->ndim() - 1)}}); @@ -78,7 +104,9 @@ std::tuple PagedAttentionImpl::do_kv_cac v_cache_layer, key, value, - slot_mapping); + slot_mapping, + k_scale, + v_scale); return {k_cache_layer, v_cache_layer}; } diff --git a/csrc/layers/attention/backends/paged_attn.hpp b/csrc/layers/attention/backends/paged_attn.hpp index 4f53ea573..e066a57e3 100644 --- a/csrc/layers/attention/backends/paged_attn.hpp +++ b/csrc/layers/attention/backends/paged_attn.hpp @@ -2,6 +2,7 @@ #include "../../../global_state/global_state.hpp" #include "infinicore/tensor.hpp" +#include #include namespace infinilm::layers::attention { @@ -40,7 +41,9 @@ class PagedAttentionImpl { const infinicore::Tensor key, const infinicore::Tensor value, infinicore::Tensor &kv_cache, - const infinicore::Tensor slot_mapping) const; + const infinicore::Tensor slot_mapping, + const std::optional &k_scale, + const std::optional &v_scale) const; private: size_t num_heads_; diff --git a/csrc/layers/quantization/fp8_blockwise.cpp b/csrc/layers/quantization/fp8_blockwise.cpp new file mode 100644 index 000000000..ead3e231f --- /dev/null +++ b/csrc/layers/quantization/fp8_blockwise.cpp @@ -0,0 +1,261 @@ +#include "fp8_blockwise.hpp" + +#include "gptq_marlin.hpp" +#include "marlin_support.hpp" +#include "marlin_utils.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace infinilm::quantization { + +namespace { + +// float -> IEEE half bits, round-to-nearest-even (scale values are finite and +// non-negative, but handle subnormals/overflow for robustness). +uint16_t f32_to_f16_bits(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof(u)); + const uint32_t sign = (u >> 16) & 0x8000u; + const int exp = static_cast((u >> 23) & 0xFFu) - 127 + 15; + const uint32_t mantissa = u & 0x7FFFFFu; + if (exp <= 0) { + if (exp < -10) { + return static_cast(sign); // underflow to zero + } + const uint32_t m = (mantissa | 0x800000u) >> (14 - exp); + return static_cast(sign | ((m + 1) >> 1)); + } + if (exp >= 31) { + return static_cast(sign | 0x7C00u); // overflow to inf + } + const uint32_t rounding = 0xFFFu + ((mantissa >> 13) & 1u); + return static_cast(sign | ((static_cast(exp) << 10) + ((mantissa + rounding) >> 13))); +} + +} // namespace + +FP8Blockwise::FP8Blockwise(const nlohmann::json &quant_config) + : BaseQuantization(quant_config) { + auto block_size = get_or>("weight_block_size", {128, 128}); + if (block_size.size() != 2 || block_size[0] == 0 || block_size[1] == 0) { + throw std::runtime_error("FP8Blockwise: weight_block_size must be [BM, BN] with positive entries"); + } + block_m_ = block_size[0]; + block_n_ = block_size[1]; +} + +std::vector FP8Blockwise::get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int /*tp_num_heads*/, + const infinicore::DataType &dtype, + bool bias) const { + if (in_features % block_n_ != 0 || out_features % block_m_ != 0) { + throw std::runtime_error("FP8Blockwise: in_features (" + std::to_string(in_features) + ") and out_features (" + std::to_string(out_features) + ") must be divisible by weight_block_size [" + std::to_string(block_m_) + ", " + std::to_string(block_n_) + "]"); + } + std::vector descs; + descs.push_back({"weight", {out_features, in_features}, infinicore::DataType::F8, split_dim, tp_rank, tp_size}); + // Scale is split across TP ranks along the same dim as the weight; the + // per-rank shard size follows from the weight shard divided by the block. + descs.push_back({"weight_scale_inv", {out_features / block_m_, in_features / block_n_}, infinicore::DataType::F32, split_dim, tp_rank, tp_size}); + if (bias) { + descs.push_back({"bias", {out_features}, dtype, split_dim >= 0 ? 0 : -1, split_dim >= 0 ? tp_rank : 0, split_dim >= 0 ? tp_size : 1}); + } + return descs; +} + +infinicore::Tensor FP8Blockwise::forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha) const { + auto input_contiguous = input->is_contiguous() ? input : input->contiguous(); + + // Fused path (decode): the fp8_blockwise_gemm kernel reads every FP8 + // weight byte exactly once and dequantizes in registers, avoiding the + // materialized BF16 weight of the naive path. The operator dispatches + // internally: SIMT warp-per-row for M <= 8 (measured on RTX 5090, + // Qwen3-8B-FP8: bs=1 141 vs 28 tok/s, bs=8 412 vs 218 against naive), + // a tensor-core mma.m16n8k16 kernel for 9 <= M <= 32 with F16/BF16 + // activations (the SIMT M_TILE kernels go instruction-throughput bound + // there; INFINIOP_FP8_GEMM_MMA=0 forces SIMT for A/B). M > 32 (prefill) + // stays on the naive cuBLAS path, as does any M with F32 activations + // (SIMT M_TILE fallback inside the operator). The fused kernel + // implements alpha == 1 only. Set INFINILM_FP8_FUSED_GEMM=0 to force + // the naive path (A/B testing). + static const bool fused_gemm_enabled = [] { + const char *env = std::getenv("INFINILM_FP8_FUSED_GEMM"); + return env == nullptr || env[0] != '0'; + }(); + const auto &weight = params.at("weight"); + const size_t k = weight->size(1); + const size_t m = input_contiguous->numel() / k; + if (fused_gemm_enabled && alpha == 1.0f && m <= 32 + && input->device().getType() == infinicore::Device::Type::NVIDIA + && block_m_ % 16 == 0 && block_n_ % 128 == 0 && k % 128 == 0) { + auto flat_input = input_contiguous->view({m, k}); + auto output = infinicore::op::fp8_blockwise_gemm( + flat_input, weight, params.at("weight_scale_inv")); + if (has_bias) { + auto bias = params.at("bias"); + infinicore::op::add_(output, output, bias->as_strided(output->shape(), {0, 1})); + } + auto out_shape = input_contiguous->shape(); + out_shape.back() = output->size(1); + return output->view(out_shape); + } + + // Naive path (prefill / fallback): dequantize the block-scaled FP8 weight + // to the activation dtype, then run the standard GEMM. + auto dequant_weight = infinicore::op::fp8_blockwise_dequantize( + weight, params.at("weight_scale_inv"), input->dtype()); + + std::optional bias_opt; + if (has_bias) { + bias_opt = params.at("bias"); + } + return infinicore::op::linear(input_contiguous, dequant_weight, bias_opt, alpha); +} + +std::vector FP8Blockwise::split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int /*tp_num_heads*/) const { + std::vector result; + const auto &weight = params.at("weight"); + const auto &weight_scale_inv = params.at("weight_scale_inv"); + const auto bias_it = params.find("bias"); + + // SplitInfo start/size are in units of the weight's narrow_dim extent. + // weight_scale_inv holds one entry per block, so its narrow range along + // the same dim is divided by the block size (BM for dim0, BN for dim1). + const size_t block = (narrow_dim == 0) ? block_m_ : block_n_; + + for (const auto &split : splits) { + if (split.start % block != 0 || split.size % block != 0) { + throw std::runtime_error("FP8Blockwise: split range at start=" + std::to_string(split.start) + " size=" + std::to_string(split.size) + " is not aligned to block size " + std::to_string(block)); + } + result.push_back({split.prefix + ".weight", + infinicore::nn::Parameter( + weight->narrow({{static_cast(narrow_dim), split.start, split.size}}), + narrow_dim, tp_rank, tp_size, split.num_shards)}); + result.push_back({split.prefix + ".weight_scale_inv", + infinicore::nn::Parameter( + weight_scale_inv->narrow({{static_cast(narrow_dim), split.start / block, split.size / block}}), + narrow_dim, tp_rank, tp_size, split.num_shards)}); + if (bias_it != params.end()) { + result.push_back({split.prefix + ".bias", + infinicore::nn::Parameter( + bias_it->second->narrow({{0, split.start, split.size}}), + 0, tp_rank, tp_size, split.num_shards)}); + } + } + return result; +} + +std::shared_ptr FP8Blockwise::process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &device, + int) const { + const auto weight_it = params.find("weight"); + const auto scale_it = params.find("weight_scale_inv"); + if (weight_it == params.end() || scale_it == params.end()) { + throw std::runtime_error( + "FP8Blockwise: post-load processing requires weight and weight_scale_inv"); + } +#if INFINILM_ENABLE_MARLIN + // Fused path (opt-in via INFINILM_FP8_MARLIN=1): convert the block-scaled + // FP8 weight to Marlin's 8-bit layout so forward runs the fused + // dequantize-in-GEMM kernel instead of materializing a BF16 weight per + // step. DeepSeek/Qwen block scales [N/BM, K/BN] are expanded (exact + // repetition) to Marlin's per-K-group, per-column grid [K/128, N]; only + // the canonical 128x128 block maps onto Marlin's group_size=128 + // (group_blocks=8) FP8 instantiation. + // + // NOTE: the InfiniCore Marlin GEMM operator requires TVM-FFI headers at + // build time (ENABLE_TVM_API; otherwise calculate() is a silent no-op), + // and its kernels are currently broken on sm_120 (deadlock on CUDA 12.8, + // garbage output on CUDA 13.2 — see dev_fp8/marlin_sm120_issue.md). The + // naive dequantize+GEMM path therefore stays the default. + static const bool fp8_marlin_enabled = [] { + const char *env = std::getenv("INFINILM_FP8_MARLIN"); + return env != nullptr && env[0] == '1'; + }(); + if (fp8_marlin_enabled && device.getType() == infinicore::Device::Type::NVIDIA && block_m_ == 128 && block_n_ == 128) { + const auto &weight = weight_it->second; // [N, K] F8 + const auto &scales = scale_it->second; // [N/128, K/128] F32 + const size_t size_n = weight->size(0); + const size_t size_k = weight->size(1); + const size_t num_groups = size_k / 128; + if (marlin::supports_shape(size_k, size_n, 128)) { + // 1) Weight: [N, K] FP8 bytes -> GPTQ-style qweight [K/4, N] I32 + // (4 consecutive K bytes per word, little-endian; gptq_value() + // extracts byte (k%4) from bits 8*(k%4)). + auto weight_cpu = weight->contiguous()->to(infinicore::Device::cpu()); + const auto *w_bytes = reinterpret_cast(weight_cpu->data()); + std::vector packed(size_k / 4 * size_n, 0); + for (size_t n = 0; n < size_n; ++n) { + const uint8_t *row = w_bytes + n * size_k; + uint32_t *dst_col = reinterpret_cast(packed.data()) + n; + for (size_t kp = 0; kp < size_k / 4; ++kp) { + dst_col[kp * size_n] = static_cast(row[4 * kp]) | (static_cast(row[4 * kp + 1]) << 8) | (static_cast(row[4 * kp + 2]) << 16) | (static_cast(row[4 * kp + 3]) << 24); + } + } + auto qweight_cpu = marlin::make_i32_tensor(packed, {size_k / 4, size_n}, infinicore::Device::cpu()); + auto perm_empty_cpu = marlin::make_empty_i32(infinicore::Device::cpu()); + // CPU input keeps the repack on the host; result is moved to the device. + params["qweight"] = marlin::gptq_marlin_repack(qweight_cpu, perm_empty_cpu, size_k, size_n, 8)->to(device); + + // 2) Scales: [N/128, K/128] F32 -> [K/128, N] with each block value + // repeated across the 128 output columns of its block (exact), + // cast to the activation dtype expected by the kernel (the bias, + // when present, is stored in the model dtype). + auto scales_cpu = scales->contiguous()->to(infinicore::Device::cpu()); + const auto *s_src = reinterpret_cast(scales_cpu->data()); + const auto bias_it = params.find("bias"); + const auto act_dtype = (bias_it != params.end()) ? bias_it->second->dtype() : infinicore::DataType::BF16; + if (act_dtype != infinicore::DataType::BF16 && act_dtype != infinicore::DataType::F16) { + return nullptr; + } + auto scales_exp = infinicore::Tensor::empty({num_groups, size_n}, act_dtype, infinicore::Device::cpu()); + auto *s_dst = reinterpret_cast(scales_exp->data()); + for (size_t g = 0; g < num_groups; ++g) { + for (size_t n = 0; n < size_n; ++n) { + const float s = s_src[(n / 128) * num_groups + g]; + uint32_t u; + std::memcpy(&u, &s, sizeof(u)); + if (act_dtype == infinicore::DataType::BF16) { + const uint32_t rounding = 0x7FFFu + ((u >> 16) & 1u); + s_dst[g * size_n + n] = static_cast((u + rounding) >> 16); + } else { + s_dst[g * size_n + n] = f32_to_f16_bits(s); + } + } + } + params["scales"] = marlin::permute_scales(scales_exp, size_k, size_n, 128)->to(device); + + params["qzeros"] = marlin::make_empty_i32(device); + params["g_idx"] = marlin::make_empty_i32(device); + params["perm"] = marlin::make_empty_i32(device); + params["global_scales"] = marlin::make_empty_i32(device); + params.erase("weight"); + params.erase("weight_scale_inv"); + + return std::make_shared( + get_config(), size_k, size_n, /*is_k_full=*/true, marlin::FE4M3FN_ID); + } + } +#endif + return nullptr; +} + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/fp8_blockwise.hpp b/csrc/layers/quantization/fp8_blockwise.hpp new file mode 100644 index 000000000..0917a7917 --- /dev/null +++ b/csrc/layers/quantization/fp8_blockwise.hpp @@ -0,0 +1,49 @@ +#pragma once + +#include "base_quantization.hpp" + +namespace infinilm::quantization { + +// FP8 (E4M3) weights with block-wise F32 scales (DeepSeek-V3 / Qwen3-FP8 +// style, quant_method == "fp8"). Weight is stored as [out, in] F8 with one +// scale per [BM, BN] block in weight_scale_inv [out / BM, in / BN] F32. +class FP8Blockwise : public BaseQuantization { +public: + explicit FP8Blockwise(const nlohmann::json &quant_config); + + QuantScheme get_quant_scheme() const override { + return QuantScheme::FP8_W8A16; + } + + std::vector get_param_layout( + size_t in_features, size_t out_features, + int split_dim, int tp_rank, int tp_size, + int tp_num_heads, + const infinicore::DataType &dtype, + bool bias) const override; + + infinicore::Tensor forward( + const ParamsMap ¶ms, + const infinicore::Tensor &input, + bool has_bias, + float alpha = 1.0f) const override; + + std::vector split_params( + const std::unordered_map ¶ms, + const std::vector &splits, + int narrow_dim, + int tp_rank, int tp_size, int tp_num_heads) const override; + + std::shared_ptr process_weights_after_loading( + ParamsMap ¶ms, + const infinicore::Device &device, + int split_dim = -1) const override; + +private: + // Block shape parsed from "weight_block_size" (default [128, 128]). + // block_m_ divides the output (dim0), block_n_ the input (dim1) extent. + size_t block_m_ = 128; + size_t block_n_ = 128; +}; + +} // namespace infinilm::quantization diff --git a/csrc/layers/quantization/gptq_marlin.cpp b/csrc/layers/quantization/gptq_marlin.cpp index e95d85b5a..7f519a0ab 100644 --- a/csrc/layers/quantization/gptq_marlin.cpp +++ b/csrc/layers/quantization/gptq_marlin.cpp @@ -72,7 +72,7 @@ infinicore::Tensor GPTQMarlin::forward( qzeros, g_idx, perm, - marlin::UINT4B8_ID, + b_q_type_id_, is_k_full_, false, true, diff --git a/csrc/layers/quantization/gptq_marlin.hpp b/csrc/layers/quantization/gptq_marlin.hpp index c16c79438..e2b413216 100644 --- a/csrc/layers/quantization/gptq_marlin.hpp +++ b/csrc/layers/quantization/gptq_marlin.hpp @@ -1,17 +1,20 @@ #pragma once #include "base_quantization.hpp" +#include "marlin_utils.hpp" namespace infinilm::quantization { class GPTQMarlin : public BaseQuantization { public: GPTQMarlin(const nlohmann::json &quant_config, size_t input_size_per_partition, - size_t output_size_per_partition, bool is_k_full) + size_t output_size_per_partition, bool is_k_full, + int64_t b_q_type_id = marlin::UINT4B8_ID) : BaseQuantization(quant_config), input_size_per_partition_(input_size_per_partition), output_size_per_partition_(output_size_per_partition), - is_k_full_(is_k_full) {} + is_k_full_(is_k_full), + b_q_type_id_(b_q_type_id) {} QuantScheme get_quant_scheme() const override { return QuantScheme::GPTQ_MARLIN_W4A16; } @@ -46,6 +49,7 @@ class GPTQMarlin : public BaseQuantization { size_t input_size_per_partition_; size_t output_size_per_partition_; bool is_k_full_; + int64_t b_q_type_id_; // Per-layer Marlin workspace. It must be all-zero before each launch // because the current InfiniCore Marlin kernels use it as lock state. // TODO: replace per-layer memset with a shared global zero workspace, or diff --git a/csrc/layers/quantization/marlin_utils.hpp b/csrc/layers/quantization/marlin_utils.hpp index edeca9aa5..88539b662 100644 --- a/csrc/layers/quantization/marlin_utils.hpp +++ b/csrc/layers/quantization/marlin_utils.hpp @@ -9,6 +9,9 @@ namespace infinilm::quantization::marlin { constexpr int64_t UINT4_ID = 1125899906843648LL; constexpr int64_t UINT4B8_ID = 1125899907892224LL; +// host::kFE4M3fn.id() from InfiniCore's sgl_kernel/scalar_type.hpp +// (FP8 E4M3 weight type for the Marlin GEMM kernel). +constexpr int64_t FE4M3FN_ID = 2814749767172868LL; bool supports_shape(size_t input_size_per_partition, size_t output_size_per_partition, int group_size); diff --git a/csrc/layers/quantization/quantization.hpp b/csrc/layers/quantization/quantization.hpp index 0cc9cd7e2..045bef98a 100644 --- a/csrc/layers/quantization/quantization.hpp +++ b/csrc/layers/quantization/quantization.hpp @@ -4,6 +4,7 @@ #include "awq_marlin.hpp" #include "base_quantization.hpp" #include "compressed_tensors.hpp" +#include "fp8_blockwise.hpp" #include "gptq.hpp" #include "gptq_marlin.hpp" #include "gptq_qy.hpp" diff --git a/csrc/layers/quantization/quantization_scheme.hpp b/csrc/layers/quantization/quantization_scheme.hpp index 455968a7a..cddde9ea6 100644 --- a/csrc/layers/quantization/quantization_scheme.hpp +++ b/csrc/layers/quantization/quantization_scheme.hpp @@ -11,11 +11,13 @@ enum class QuantScheme { GPTQ_W4A16, GPTQ_MARLIN_W4A16, MXFP4_W4A16, + FP8_W8A16, }; enum class KVQuantAlgo { NONE, INT8, + FP8, }; } // namespace infinilm::quantization diff --git a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp index 3c8a3f8b6..0bd6589d0 100644 --- a/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp +++ b/csrc/models/deepseek_v2/deepseek_v2_mla_attention.cpp @@ -193,7 +193,7 @@ infinicore::Tensor DeepseekV2MLAAttention::forward_paged_(const infinicore::Tens auto k_cache_layer = kv_cache->narrow({{3, 0, mla_head_dim_}}); auto v_cache_layer = kv_cache->narrow({{3, mla_head_dim_, kv_lora_rank_}}); - infinicore::op::paged_caching_(k_cache_layer, v_cache_layer, key_states, value_states, slot_mapping.value()); + infinicore::op::paged_caching_(k_cache_layer, v_cache_layer, key_states, value_states, slot_mapping.value(), std::nullopt, std::nullopt); auto attn_output = infinicore::Tensor::empty({seq_len, num_attention_heads_, kv_lora_rank_}, query_states->dtype(), query_states->device()); const bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]); @@ -230,7 +230,9 @@ infinicore::Tensor DeepseekV2MLAAttention::forward_paged_(const infinicore::Tens total_sequence_lengths.value(), input_offsets.value(), std::nullopt, - softmax_scale_); + softmax_scale_, + std::nullopt, + std::nullopt); } } else { infinicore::op::paged_attention_( @@ -241,7 +243,9 @@ infinicore::Tensor DeepseekV2MLAAttention::forward_paged_(const infinicore::Tens block_tables.value(), total_sequence_lengths.value(), std::nullopt, - softmax_scale_); + softmax_scale_, + std::nullopt, + std::nullopt); } return project_latent_to_value_(attn_output, batch_size, seq_len); } diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index 5d284a316..f8edd075e 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -10,6 +10,7 @@ void InfinilmModel::reset_cache(const cache::CacheConfig *cache_config) { if (cache_config == nullptr) { cache_config_.reset(); global_state::get_forward_context().kv_cache_vec.clear(); + global_state::get_forward_context().kv_scale_vec.clear(); return; } cache_config_ = cache_config->unique_copy(); @@ -33,6 +34,9 @@ std::vector InfinilmModel::default_allocate_kv_cache_tensors size_t num_key_value_heads = text_config->get("num_key_value_heads"); size_t max_position_embeddings = text_config->get("max_position_embeddings"); const auto &dtype = model_config_->get_kv_cache_dtype(); + if (dtype == infinicore::DataType::F8 && attention_backend != backends::AttentionBackend::PAGED_ATTN) { + throw std::runtime_error("infinilm::InfinilmModel::default_allocate_kv_cache_tensors: FP8 KV cache is only supported on the PAGED_ATTN backend"); + } const size_t num_hidden_layers = text_config->get("num_hidden_layers"); const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); const size_t pp_size = static_cast(rank_info.pp_size); @@ -73,6 +77,16 @@ std::vector InfinilmModel::default_allocate_kv_cache_tensors } kv_cache_vec.resize(num_hidden_layers); + // FP8(E4M3) KV cache: allocate per-layer k_scale/v_scale alongside the cache. + // They live in ForwardContext::kv_scale_vec, parallel to kv_cache_vec, and are + // consumed by PagedAttentionImpl. Cleared here so a non-FP8 reset drops stale scales. + const bool kv_fp8 = (dtype == infinicore::DataType::F8); + auto &kv_scale_vec = global_state::get_forward_context().kv_scale_vec; + kv_scale_vec.clear(); + if (kv_fp8) { + kv_scale_vec.resize(num_hidden_layers); + } + for (size_t layer_idx = local_layer_begin; layer_idx < local_layer_end; ++layer_idx) { auto kv_cache = cache::PagedKVCache::create_layer_kv_cache( head_dim, @@ -82,6 +96,11 @@ std::vector InfinilmModel::default_allocate_kv_cache_tensors dtype, *paged_kv_cache_config); kv_cache_vec[layer_idx] = kv_cache; + if (kv_fp8) { + kv_scale_vec[layer_idx] = cache::PagedKVCache::create_layer_kv_scales( + num_key_value_heads, + *paged_kv_cache_config); + } } infinicore::context::syncStream(); break; diff --git a/csrc/models/videonsa/videonsa_attention.cpp b/csrc/models/videonsa/videonsa_attention.cpp index 828d10860..4574ddac3 100644 --- a/csrc/models/videonsa/videonsa_attention.cpp +++ b/csrc/models/videonsa/videonsa_attention.cpp @@ -380,7 +380,7 @@ infinicore::Tensor VideoNSAAttention::forward(const infinicore::Tensor &position auto k_cache_for_nsa = is_flash_attn ? k_cache_layer->permute({0, 2, 1, 3}) : k_cache_layer; auto v_cache_for_nsa = is_flash_attn ? v_cache_layer->permute({0, 2, 1, 3}) : v_cache_layer; - infinicore::op::paged_caching_(k_cache_for_nsa, v_cache_for_nsa, k_reshaped, v_reshaped, attn_metadata.slot_mapping.value()); + infinicore::op::paged_caching_(k_cache_for_nsa, v_cache_for_nsa, k_reshaped, v_reshaped, attn_metadata.slot_mapping.value(), std::nullopt, std::nullopt); if (can_use_paged_decode_nsa) { auto gate_hidden = g_proj_1_->forward(hidden_states_mutable); diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c0..fe3056fef 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -287,7 +287,7 @@ def _add_common_args(self): "--kv-cache-dtype", type=str, default=None, - choices=["int8"], + choices=["int8", "fp8", "float8"], help="KV cache data type", ) self.parser.add_argument( diff --git a/python/infinilm/config/engine_config.py b/python/infinilm/config/engine_config.py index bccb7e758..e05fb8709 100644 --- a/python/infinilm/config/engine_config.py +++ b/python/infinilm/config/engine_config.py @@ -27,6 +27,8 @@ class EngineConfig: num_blocks: Number of KV cache blocks (only for paged cache). block_size: Size of each KV cache block (only for paged cache). max_cache_len: Maximum sequence length (only for static cache). + kv_cache_dtype: KV cache data type ('int8', 'fp8'); None keeps the model dtype. + 'fp8' requires the paged attention backend (attn_backend='paged-attn'). enable_prefix_caching: Whether to reuse KV cache across requests. temperature: Default sampling temperature. top_p: Default top-p sampling parameter. @@ -57,6 +59,7 @@ class EngineConfig: num_blocks: int = 512 block_size: int = 256 max_cache_len: int = 4096 + kv_cache_dtype: Optional[str] = None temperature: float = 1.0 top_p: float = 0.8 top_k: int = 1 diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..37ee5f4e7 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -360,6 +360,7 @@ def __init__( top_k: int = 1, enable_graph: bool = False, attn_backend: str = "default", + kv_cache_dtype: Optional[str] = None, use_mla: bool = False, pre_transpose: bool = False, weight_load_mode: str = "async", @@ -385,6 +386,8 @@ def __init__( top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. attn_backend: Attention backend to use ('default', 'flash-attn'). + kv_cache_dtype: KV cache data type ('int8', 'fp8'); None keeps the model + dtype. 'fp8' requires attn_backend='paged-attn'. use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. """ @@ -412,6 +415,7 @@ def __init__( top_k=top_k, enable_graph=enable_graph, attn_backend=attn_backend, + kv_cache_dtype=kv_cache_dtype, use_mla=use_mla, pre_transpose=pre_transpose, weight_load_mode=weight_load_mode, @@ -587,6 +591,7 @@ def __init__( top_k: int = 1, enable_graph: bool = False, attn_backend: str = "default", + kv_cache_dtype: Optional[str] = None, kv_transfer_config: Optional[KVTransferConfig] = None, use_mla: bool = False, pre_transpose: bool = False, @@ -613,6 +618,8 @@ def __init__( top_k: Default top-k sampling parameter. enable_graph: Whether to enable graph compiling. attn_backend: Attention backend to use ('default', 'flash-attn'). + kv_cache_dtype: KV cache data type ('int8', 'fp8'); None keeps the model + dtype. 'fp8' requires attn_backend='paged-attn'. kv_connector: KV connector type ('MooncakeConnector'). kv_role: Role in KV connector ('kv_producer' or 'kv_consumer'). kv_connector_extra_config: Extra config dict for KV connector. @@ -644,6 +651,7 @@ def __init__( top_k=top_k, enable_graph=enable_graph, attn_backend=attn_backend, + kv_cache_dtype=kv_cache_dtype, kv_transfer_config=kv_transfer_config, use_mla=use_mla, pre_transpose=pre_transpose, diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index a1696f848..ad6a15864 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -87,6 +87,7 @@ def __init__(self, config: EngineConfig, initialize_processor: bool = True): cache_config=cache_config, enable_graph_compiling=config.enable_graph, attention_backend=config.attn_backend, + kv_cache_dtype=config.kv_cache_dtype, use_mla=config.use_mla, weight_load_mode=config.weight_load_mode, use_legacy_moe=config.use_legacy_moe, diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..d5adc6690 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -30,6 +30,8 @@ def parse_dtype(dtype_str: str): return infinicore.float16 elif dtype_str == "bfloat16": return infinicore.bfloat16 + elif dtype_str in ("fp8", "float8"): + return infinicore.float8 elif dtype_str == "int8": return infinicore.int8 elif dtype_str == "int32": @@ -105,9 +107,15 @@ def load_state_dict( device="cpu", dtype=torch.bfloat16, preserve_fp32_suffixes: Tuple[str, ...] = (".e_score_correction_bias",), + preserve_dtype_suffixes: Tuple[str, ...] = (), + preserve_dtypes: Tuple[torch.dtype, ...] = (), ) -> Dict[str, torch.Tensor]: """ Reads a `safetensor` checkpoint file. We load the checkpoint on "cpu" by default. + + Tensors keep their checkpoint dtype when their key ends with one of + `preserve_dtype_suffixes` or their dtype is listed in `preserve_dtypes` + (e.g. float8 storage formats that must not be silently dequantized). """ if not checkpoint_file.endswith(".safetensors"): @@ -129,7 +137,10 @@ def load_state_dict( for k in f.keys(): tensor = f.get_tensor(k) preserve_fp32 = k.endswith(preserve_fp32_suffixes) - if tensor.is_floating_point() and not preserve_fp32: + preserve_dtype = k.endswith(preserve_dtype_suffixes) or ( + tensor.dtype in preserve_dtypes + ) + if tensor.is_floating_point() and not (preserve_fp32 or preserve_dtype): tensor = tensor.to(device=device, dtype=dtype) else: tensor = tensor.to(device=device) @@ -189,6 +200,57 @@ def get_model_state_dict( return model_param_infini +def _resolve_preserve_config( + hf_config: dict, +) -> Tuple[Tuple[str, ...], Tuple[str, ...], Tuple[torch.dtype, ...]]: + """Compute (preserve_fp32_suffixes, preserve_dtype_suffixes, preserve_dtypes) + for a model from its HF config. + + Tensors matching these keep their checkpoint dtype (or fp32) instead of + being cast to the model compute dtype during loading. + """ + model_type = hf_config.get("model_type", "") + preserve_fp32_suffixes = (".e_score_correction_bias",) + preserve_dtype_suffixes = () + preserve_dtypes = () + if model_type == "kimi_k3": + preserve_fp32_suffixes += (".A_log", ".dt_bias") + + # FP8 block-quantized checkpoints (Qwen3-FP8 / DeepSeek-V3 style): weights + # are stored as F8_E4M3 and block scales as F32 ``*weight_scale_inv``. + # Neither may go through the default float cast -- it would silently + # dequantize the weights and truncate scale precision before the + # quantization scheme gets to consume them. + quant_config = hf_config.get("quantization_config") or {} + if quant_config.get("quant_method") == "fp8": + preserve_dtypes += (torch.float8_e4m3fn, torch.float8_e5m2) + preserve_dtype_suffixes += (".weight_scale_inv", ".weight_scale") + + return preserve_fp32_suffixes, preserve_dtype_suffixes, preserve_dtypes + + +def _cast_fp8_scales_to_fp32( + model_param: Dict[str, torch.Tensor], hf_config: dict +) -> Dict[str, torch.Tensor]: + """Widen fp8 block-quantization scales to float32. + + FP8 checkpoints (``quant_method == "fp8"``) may store + ``*.weight_scale_inv`` / ``*.weight_scale`` in reduced precision + (Qwen3-8B-FP8 stores them as BF16), while the C++ + ``fp8_blockwise_dequantize`` operator only accepts F32 scales. + Widening is lossless. No-op for non-fp8 configs. + """ + quant_config = hf_config.get("quantization_config") or {} + if quant_config.get("quant_method") != "fp8": + return model_param + for key, tensor in model_param.items(): + if key.endswith((".weight_scale_inv", ".weight_scale")) and ( + tensor.dtype != torch.float32 + ): + model_param[key] = tensor.to(torch.float32) + return model_param + + def load_model_state_dict_by_file( model: infinicore.nn.Module, model_path: str, @@ -201,9 +263,9 @@ def load_model_state_dict_by_file( t1 = time.time() model_type = model.hf_config.get("model_type", "") - preserve_fp32_suffixes = (".e_score_correction_bias",) - if model_type == "kimi_k3": - preserve_fp32_suffixes += (".A_log", ".dt_bias") + preserve_fp32_suffixes, preserve_dtype_suffixes, preserve_dtypes = ( + _resolve_preserve_config(model.hf_config) + ) torch_device = "cpu" torch_dtype = infinicore.utils.to_torch_dtype(dtype) @@ -246,11 +308,14 @@ def load_model_state_dict_by_file( device=torch_device, dtype=torch_dtype, preserve_fp32_suffixes=preserve_fp32_suffixes, + preserve_dtype_suffixes=preserve_dtype_suffixes, + preserve_dtypes=preserve_dtypes, ) # Apply model-specific weight remapping if remapper is not None: model_param = remapper(model_param, config=model.hf_config) + model_param = _cast_fp8_scales_to_fp32(model_param, model.hf_config) # --------------------------------------------------------- # # Scale embed_tokens on torch side before converting @@ -301,6 +366,7 @@ def load_model_state_dict_by_file( remapper = _WEIGHT_REMAPPER.get(model_type) if remapper is not None: model_params = remapper(model_params, config=model.hf_config) + model_params = _cast_fp8_scales_to_fp32(model_params, model.hf_config) # Scale embed_tokens on torch side before converting if "model.embed_tokens.weight" in model_params: diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..997bfba3e 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -118,6 +118,7 @@ def __init__( port: int = 8000, enable_graph: bool = False, attn_backend: str = "default", + kv_cache_dtype: Optional[str] = None, use_mla: bool = False, skip_load: bool = False, weight_load_mode: str = "async", @@ -149,6 +150,8 @@ def __init__( port: Server port number. enable_graph: Whether to enable graph compiling. attn_backend: Attention backend to use ('default', 'flash-attn'). + kv_cache_dtype: KV cache data type ('int8', 'fp8'); None keeps the model + dtype. 'fp8' requires attn_backend='paged-attn'. use_mla: Whether to use DeepSeek V2 MLA attention when supported. skip_load: Whether to skip loading model weights. weight_load_mode: Weight loading mode across tensor-parallel workers. @@ -181,6 +184,7 @@ def __init__( self.port = port self.enable_graph = enable_graph self.attn_backend = attn_backend + self.kv_cache_dtype = kv_cache_dtype self.use_mla = use_mla self.skip_load = skip_load self.weight_load_mode = weight_load_mode @@ -226,6 +230,7 @@ async def lifespan(app: FastAPI): top_k=self.top_k, enable_graph=self.enable_graph, attn_backend=self.attn_backend, + kv_cache_dtype=self.kv_cache_dtype, use_mla=self.use_mla, skip_load=self.skip_load, weight_load_mode=self.weight_load_mode, @@ -660,6 +665,7 @@ def main(): port=cfg.port, enable_graph=cfg.enable_graph, attn_backend=cfg.attn, + kv_cache_dtype=cfg.kv_cache_dtype, use_mla=cfg.use_mla, skip_load=cfg.skip_load, weight_load_mode=cfg.weight_load_mode, diff --git a/python/infinilm/server/pipeline_worker.py b/python/infinilm/server/pipeline_worker.py index 9e4bb8e6b..738e6a9b2 100644 --- a/python/infinilm/server/pipeline_worker.py +++ b/python/infinilm/server/pipeline_worker.py @@ -34,6 +34,7 @@ def run_worker(cfg: BaseConfig) -> None: top_k=cfg.top_k, enable_graph=cfg.enable_graph, attn_backend=cfg.attn, + kv_cache_dtype=cfg.kv_cache_dtype, use_mla=cfg.use_mla, weight_load_mode=cfg.weight_load_mode, skip_load=cfg.skip_load, From b9897c674397998066030823c7338e56aca5290b Mon Sep 17 00:00:00 2001 From: yangyeyuan <2776079516@qq.com> Date: Sun, 20 Sep 2026 15:25:17 +0800 Subject: [PATCH 2/4] test: add Qwen3 FP8 checkpoint loading test --- .../qwen3_fp8/test_fp8_checkpoint_loading.py | 165 ++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 test/models/qwen3_fp8/test_fp8_checkpoint_loading.py diff --git a/test/models/qwen3_fp8/test_fp8_checkpoint_loading.py b/test/models/qwen3_fp8/test_fp8_checkpoint_loading.py new file mode 100644 index 000000000..584d0ef60 --- /dev/null +++ b/test/models/qwen3_fp8/test_fp8_checkpoint_loading.py @@ -0,0 +1,165 @@ +import os +import tempfile +import unittest + +import torch +from infinilm.modeling_utils import ( + _cast_fp8_scales_to_fp32, + _resolve_preserve_config, + load_state_dict, +) +from safetensors.torch import save_file + +QWEN3_FP8_HF_CONFIG = { + "model_type": "qwen3", + "torch_dtype": "bfloat16", + "quantization_config": { + "quant_method": "fp8", + "fmt": "e4m3", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + }, +} + + +class ResolvePreserveConfigTest(unittest.TestCase): + def test_fp8_quant_config_preserves_float8_and_scales(self): + fp32_suffixes, dtype_suffixes, dtypes = _resolve_preserve_config( + QWEN3_FP8_HF_CONFIG + ) + + self.assertIn(torch.float8_e4m3fn, dtypes) + self.assertIn(torch.float8_e5m2, dtypes) + self.assertIn(".weight_scale_inv", dtype_suffixes) + self.assertIn(".weight_scale", dtype_suffixes) + self.assertEqual(fp32_suffixes, (".e_score_correction_bias",)) + + def test_plain_config_keeps_default_behavior(self): + fp32_suffixes, dtype_suffixes, dtypes = _resolve_preserve_config( + {"model_type": "qwen3"} + ) + + self.assertEqual(fp32_suffixes, (".e_score_correction_bias",)) + self.assertEqual(dtype_suffixes, ()) + self.assertEqual(dtypes, ()) + + + +class FP8CheckpointLoadingTest(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.checkpoint_path = os.path.join( + self.temp_dir.name, "model-00001-of-00001.safetensors" + ) + + self.weight = torch.tensor( + [[1.0, -2.0, 0.5, 0.25], [3.0, -4.0, 5.5, -6.0]], + dtype=torch.float32, + ).to(torch.float8_e4m3fn) + self.scale_inv = torch.tensor([[0.015625]], dtype=torch.float32) + self.state_dict = { + "model.layers.0.mlp.gate_proj.weight": self.weight, + "model.layers.0.mlp.gate_proj.weight_scale_inv": self.scale_inv, + "model.layers.0.input_layernorm.weight": torch.tensor( + [0.1, 0.2], dtype=torch.float32 + ), + } + save_file(self.state_dict, self.checkpoint_path, metadata={"format": "pt"}) + + def tearDown(self): + self.temp_dir.cleanup() + + def test_fp8_config_preserves_weight_bits_and_scale_dtype(self): + fp32_suffixes, dtype_suffixes, dtypes = _resolve_preserve_config( + QWEN3_FP8_HF_CONFIG + ) + loaded = load_state_dict( + self.checkpoint_path, + dtype=torch.bfloat16, + preserve_fp32_suffixes=fp32_suffixes, + preserve_dtype_suffixes=dtype_suffixes, + preserve_dtypes=dtypes, + ) + + weight = loaded["model.layers.0.mlp.gate_proj.weight"] + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertTrue( + torch.equal(weight.view(torch.uint8), self.weight.view(torch.uint8)) + ) + + scale = loaded["model.layers.0.mlp.gate_proj.weight_scale_inv"] + self.assertEqual(scale.dtype, torch.float32) + self.assertTrue(torch.equal(scale, self.scale_inv)) + + # Non-quantized float tensors still follow the model compute dtype. + self.assertEqual( + loaded["model.layers.0.input_layernorm.weight"].dtype, torch.bfloat16 + ) + + def test_default_path_still_casts_float8_weights(self): + loaded = load_state_dict(self.checkpoint_path, dtype=torch.bfloat16) + + self.assertEqual( + loaded["model.layers.0.mlp.gate_proj.weight"].dtype, torch.bfloat16 + ) + self.assertEqual( + loaded["model.layers.0.mlp.gate_proj.weight_scale_inv"].dtype, + torch.bfloat16, + ) + + +class CastFP8ScalesToFP32Test(unittest.TestCase): + def test_bf16_scales_widened_to_fp32(self): + # Real Qwen3-8B-FP8 checkpoints store weight_scale_inv as BF16. + scale_inv = torch.tensor([[0.015625, 0.5], [-0.25, 1.0]], dtype=torch.bfloat16) + weight = torch.zeros(2, 2, dtype=torch.float32).to(torch.float8_e4m3fn) + params = { + "model.layers.0.mlp.gate_proj.weight": weight, + "model.layers.0.mlp.gate_proj.weight_scale_inv": scale_inv, + "model.layers.0.input_layernorm.weight": torch.tensor( + [0.1, 0.2], dtype=torch.bfloat16 + ), + } + + out = _cast_fp8_scales_to_fp32(params, QWEN3_FP8_HF_CONFIG) + + scale = out["model.layers.0.mlp.gate_proj.weight_scale_inv"] + self.assertEqual(scale.dtype, torch.float32) + self.assertTrue(torch.equal(scale, scale_inv.to(torch.float32))) + # FP8 weight bits and unrelated tensors are left untouched. + self.assertIs(out["model.layers.0.mlp.gate_proj.weight"], weight) + self.assertEqual( + out["model.layers.0.input_layernorm.weight"].dtype, torch.bfloat16 + ) + + def test_weight_scale_suffix_also_cast(self): + scale = torch.tensor([0.5], dtype=torch.bfloat16) + params = {"model.layers.0.self_attn.q_proj.weight_scale": scale} + + out = _cast_fp8_scales_to_fp32(params, QWEN3_FP8_HF_CONFIG) + + self.assertEqual( + out["model.layers.0.self_attn.q_proj.weight_scale"].dtype, + torch.float32, + ) + + def test_fp32_scale_kept_as_is(self): + scale = torch.tensor([[0.015625]], dtype=torch.float32) + params = {"model.layers.0.mlp.gate_proj.weight_scale_inv": scale} + + out = _cast_fp8_scales_to_fp32(params, QWEN3_FP8_HF_CONFIG) + + self.assertIs(out["model.layers.0.mlp.gate_proj.weight_scale_inv"], scale) + + def test_non_fp8_config_is_noop(self): + scale = torch.tensor([[0.015625]], dtype=torch.bfloat16) + params = {"model.layers.0.mlp.gate_proj.weight_scale_inv": scale} + + out = _cast_fp8_scales_to_fp32(params, {"model_type": "qwen3"}) + + self.assertIs(out["model.layers.0.mlp.gate_proj.weight_scale_inv"], scale) + self.assertEqual(scale.dtype, torch.bfloat16) + + +if __name__ == "__main__": + unittest.main() From 48fc9a250f9a9a0b44b684519f1d9fe7ac8d3598 Mon Sep 17 00:00:00 2001 From: yangyeyuan <2776079516@qq.com> Date: Sun, 20 Sep 2026 15:25:29 +0800 Subject: [PATCH 3/4] build: declare runtime dependencies and server/multimodal extras dependencies was empty, so a fresh environment could not import infinilm.llm (janus, xxhash missing). Core deps cover the LLM path; fastapi/uvicorn/httpx/pydantic/msgspec/pyzmq/psutil move to the server extra and pillow to multimodal. --- pyproject.toml | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3ee4742c3..5dc3fea74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,16 @@ name = "InfiniLM" version = "0.1.0" description = "InfiniLM model implementations" readme = "README.md" -dependencies = [] +dependencies = [ + "torch", + "transformers", + "numpy", + "safetensors", + "tokenizers", + "tqdm", + "janus", + "xxhash", +] requires-python = ">=3.10" classifiers = [ "Programming Language :: Python :: 3", @@ -15,6 +24,18 @@ classifiers = [ "Operating System :: OS Independent", ] +[project.optional-dependencies] +server = [ + "fastapi", + "uvicorn", + "httpx", + "pydantic", + "msgspec", + "pyzmq", + "psutil", +] +multimodal = ["pillow"] + [project.urls] Homepage = "https://github.com/InfiniTensor/InfiniLM" From 88ebee4a10d28fca88ba424b7a9e4b77d2b9536c Mon Sep 17 00:00:00 2001 From: yangyeyuan <2776079516@qq.com> Date: Sun, 20 Sep 2026 15:25:29 +0800 Subject: [PATCH 4/4] fix: log expected shutdown exceptions as warnings, not fatal The top-level thread_loop catch always logged "fatal exception" even when the worker had already been asked to stop (teardown racing an in-flight CUDA call), which made normal shutdowns indistinguishable from real faults in the logs. Expected deaths are now warnings; genuine failures remain error-level and still propagate to waiters via should_exit_. --- csrc/engine/rank_worker.cpp | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 0fa0a84cf..731b13484 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -575,14 +575,24 @@ void RankWorker::thread_loop() { compiler_.reset(); } catch (const std::exception &e) { // Top-level exception: ensure any waiters are woken and the thread exits cleanly. + bool was_exiting; { std::lock_guard lk(mutex_); + was_exiting = should_exit_; init_done_ = true; should_exit_ = true; job_done_ = true; } cv_.notify_all(); - spdlog::error("[{}] fatal exception in thread_loop: {} \n", info(), e.what()); + if (was_exiting) { + // Already asked to stop (e.g. teardown raced an in-flight job or a + // CUDA call failed while the destructor was joining the thread): + // an expected death, not a fatal fault. Keep it visible but do not + // mislabel it as fatal. + spdlog::warn("[{}] exception in thread_loop during shutdown: {} \n", info(), e.what()); + } else { + spdlog::error("[{}] fatal exception in thread_loop: {} \n", info(), e.what()); + } } }