feat: add FP8 blockwise weight quantization and FP8(E4M3) KV cache support

csrc/cache/kv_cache.cpp

@@ -146,6 +146,23 @@ infinicore::Tensor create_layer_kv_cache(
 
     return kv_cache;
 }
+
+std::pair<infinicore::Tensor, infinicore::Tensor> 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

csrc/cache/kv_cache.hpp

@@ -5,6 +5,7 @@
 
 #include <limits>
 #include <memory>
+#include <utility>
 
 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<infinicore::Tensor, infinicore::Tensor> create_layer_kv_scales(
+    infinicore::Size num_kv_heads,
+    const PagedKVCacheConfig &config);
+
 } // namespace PagedKVCache
 
 } // namespace infinilm::cache

csrc/config/quant_config.cpp

@@ -22,6 +22,8 @@ QuantConfig::get_quantization_method() const {
         return std::make_shared<infinilm::quantization::GPTQ>(quantization_config);
     } else if (quant_method == "quark") {
         return std::make_shared<infinilm::quantization::MXFP4>(quantization_config);
+    } else if (quant_method == "fp8") {
+        return std::make_shared<infinilm::quantization::FP8Blockwise>(quantization_config);
     } else {
         return std::make_shared<infinilm::quantization::NoneQuantization>(quantization_config);
     }

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;

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

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<std::mutex> 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());
+        }
     }
 }
 

csrc/global_state/forward_context.hpp

@@ -2,6 +2,8 @@
 
 #include "../models/infinilm_model.hpp"
 
+#include <utility>
+
 namespace infinilm::global_state {
 
 struct AttentionMetadata {
@@ -68,6 +70,11 @@ struct ForwardContext {
     MambaMetadata mamba_metadata;
     MultiModalMetadata mm_metadata;
     std::vector<infinicore::Tensor> 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<std::pair<infinicore::Tensor, infinicore::Tensor>> kv_scale_vec;
     std::vector<infinicore::Tensor> conv_state_vec;
     std::vector<infinicore::Tensor> ssm_state_vec;
 };

csrc/layers/attention/attention.cpp

@@ -150,6 +150,10 @@ void init_kv_cache_quant_params(std::function<void(const std::string &, infinico
         kv_cache_v_scale = infinicore::nn::Parameter({1}, infinicore::DataType::F32, device, 0, 0, 1);
         register_fn("kv_cache_v_scale", kv_cache_v_scale);
         break;
+    case infinilm::quantization::KVQuantAlgo::FP8:
+        // FP8 KV cache uses per-token-per-kv-head dynamic scales allocated together with
+        // the paged cache (ForwardContext::kv_scale_vec), not per-tensor module parameters.
+        break;
     default:
         throw std::runtime_error("unsupported kv_quant_scheme");
     }

csrc/layers/attention/backends/flash_attn.cpp

@@ -101,7 +101,9 @@ std::tuple<infinicore::Tensor, infinicore::Tensor> 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};
 }

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<infinicore::Tensor> k_scale = std::nullopt;
+    std::optional<infinicore::Tensor> 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<infinicore::Tensor, infinicore::Tensor> 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<infinicore::Tensor> &k_scale,
+                                                                                          const std::optional<infinicore::Tensor> &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<infinicore::Tensor, infinicore::Tensor> PagedAttentionImpl::do_kv_cac
         v_cache_layer,
         key,
         value,
-        slot_mapping);
+        slot_mapping,
+        k_scale,
+        v_scale);
 
     return {k_cache_layer, v_cache_layer};
 }

csrc/layers/attention/backends/paged_attn.hpp

@@ -2,6 +2,7 @@
 
 #include "../../../global_state/global_state.hpp"
 #include "infinicore/tensor.hpp"
+#include <optional>
 #include <tuple>
 
 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<infinicore::Tensor> &k_scale,
+                                                                          const std::optional<infinicore::Tensor> &v_scale) const;
 
 private:
     size_t num_heads_;