Merge branch 'demo131' into Issue/862

include/infinicore/ops/kv_caching.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(KVCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &);
+
+void kv_caching_(Tensor k_cache,
+                 Tensor v_cache,
+                 const Tensor &k,
+                 const Tensor &v,
+                 const Tensor &past_kv_lengths);
+} // namespace infinicore::op

include/infinicore/ops/linear_w4a16_awq.hpp

+#pragma once
+
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+Tensor linear_w4a16_awq(Tensor input, Tensor weight_packed, Tensor weight_scale, Tensor weight_zeros, std::optional<Tensor> bias);
+
+void linear_w4a16_awq_(Tensor out, Tensor input, Tensor weight_packed, Tensor weight_scale, Tensor weight_zeros, std::optional<Tensor> bias);
+
+} // namespace infinicore::op

include/infinicore/ops/linear_w8a8i8.hpp

+#pragma once
+
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+Tensor linear_w8a8i8(Tensor input, Tensor weight_packed, Tensor weight_scale, std::optional<Tensor> bias);
+
+void linear_w8a8i8_(Tensor out, Tensor input, Tensor weight_packed, Tensor weight_scale, std::optional<Tensor> bias);
+
+} // namespace infinicore::op

include/infinicore/ops/mul.hpp

 #pragma once
 
 #include "../device.hpp"
+#include "../graph/graph.hpp"
 #include "common/op.hpp"
 
 namespace infinicore::op {
-class Mul {
-public:
-    using schema = void (*)(Tensor, Tensor, Tensor);
-    static void execute(Tensor c, Tensor a, Tensor b);
-    static common::OpDispatcher<schema> &dispatcher();
-};
 
-Tensor mul(Tensor a, Tensor b);
-void mul_(Tensor c, Tensor a, Tensor b);
+INFINICORE_GRAPH_OP_CLASS(Mul, Tensor, const Tensor &, const Tensor &);
+
+Tensor mul(const Tensor &a, const Tensor &b);
+void mul_(Tensor c, const Tensor &a, const Tensor &b);
+
 } // namespace infinicore::op

include/infinicore/ops/paged_attention.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(PagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional<Tensor>, float);
+
+Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache,
+                       const Tensor &block_tables, const Tensor &kv_lens,
+                       std::optional<Tensor> alibi_slopes, float scale);
+
+void paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache,
+                      const Tensor &block_tables, const Tensor &kv_lens,
+                      std::optional<Tensor> alibi_slopes, float scale);
+
+} // namespace infinicore::op

include/infinicore/ops/paged_attention_prefill.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+class PagedAttentionPrefill {
+public:
+    /**
+     * @brief PagedAttentionPrefill operator signature
+     * * Argument order:
+     * 1. out: Output tensor (Packed format)
+     * 2. q: Current Query tensor (Packed format)
+     * 3. k_cache: Physical Key cache (Paged format)
+     * 4. v_cache: Physical Value cache (Paged format)
+     * 5. block_tables: Mapping table from logical blocks to physical blocks
+     * 6. total_kv_lens:  lengths of Complete Key/Value for each request
+     * 7. cu_seqlens_q: Cumulative sequence lengths of Query (prefix sum for variable-length batch)
+     * 8. alibi_slopes: ALiBi bias slopes (optional)
+     * 9. scale: Scaling factor (typically 1/sqrt(head_size))
+     */
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, std::optional<Tensor>, float);
+
+    static void execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache,
+                        Tensor block_tables, Tensor total_kv_lens, Tensor cum_seqlens_q,
+                        std::optional<Tensor> alibi_slopes, float scale);
+
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor paged_attention_prefill(Tensor q,
+                               Tensor k_cache,
+                               Tensor v_cache,
+                               Tensor block_tables,
+                               Tensor total_kv_lens,
+                               Tensor cum_seqlens_q,
+                               std::optional<Tensor> alibi_slopes,
+                               float scale);
+
+void paged_attention_prefill_(Tensor out,
+                              Tensor q,
+                              Tensor k_cache,
+                              Tensor v_cache,
+                              Tensor block_tables,
+                              Tensor total_kv_lens,
+                              Tensor cum_seqlens_q,
+                              std::optional<Tensor> alibi_slopes,
+                              float scale);
+
+} // namespace infinicore::op

include/infinicore/ops/paged_caching.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(PagedCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &);
+
+void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping);
+
+} // namespace infinicore::op

include/infinicore/ops/per_channel_quant_i8.hpp

+#pragma once
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(PerChannelQuantI8, const Tensor &, Tensor, Tensor);
+
+void per_channel_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale);
+} // namespace infinicore::op

include/infinicore/ops/rearrange.hpp

 #pragma once
 
 #include "../device.hpp"
+#include "../graph/graph.hpp"
 #include "common/op.hpp"
 
 namespace infinicore::op {
-class Rearrange {
-public:
-    using schema = void (*)(Tensor, Tensor);
-    static void execute(Tensor y, Tensor x);
-    static common::OpDispatcher<schema> &dispatcher();
-};
 
-Tensor rearrange(Tensor x);
-void rearrange_(Tensor y, Tensor x);
+INFINICORE_GRAPH_OP_CLASS(Rearrange, Tensor, const Tensor &);
+
+Tensor rearrange(const Tensor &x);
+void rearrange_(Tensor y, const Tensor &x);
+
 } // namespace infinicore::op

include/infinicore/ops/rms_norm.hpp

 #pragma once
 
 #include "../device.hpp"
+#include "../graph/graph.hpp"
 #include "common/op.hpp"
 
 namespace infinicore::op {
-class RMSNorm {
-public:
-    using schema = void (*)(Tensor, Tensor, Tensor, float);
-    static void execute(Tensor y, Tensor x, Tensor weight, float epsilon = 1e-5f);
-    static common::OpDispatcher<schema> &dispatcher();
-};
 
-Tensor rms_norm(Tensor x, Tensor weight, float epsilon = 1e-5f);
-void rms_norm_(Tensor y, Tensor x, Tensor weight, float epsilon = 1e-5f);
+INFINICORE_GRAPH_OP_CLASS(RMSNorm, Tensor, const Tensor &, const Tensor &, float);
+
+Tensor rms_norm(const Tensor &x, const Tensor &weight, float epsilon = 1e-5f);
+void rms_norm_(Tensor y, const Tensor &x, const Tensor &weight, float epsilon = 1e-5f);
+
 } // namespace infinicore::op

include/infinicore/ops/rope.hpp

 #pragma once
 
 #include "../device.hpp"
+#include "../graph/graph.hpp"
 #include "../nn/rope.hpp"
 #include "../tensor.hpp"
 #include "common/op.hpp"
 
 namespace infinicore::op {
-class RoPE {
-public:
-    using schema = void (*)(Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infinicore::nn::RoPE::Algo);
-    static void execute(Tensor x_out, const Tensor &x, const Tensor &pos, const Tensor &sin_table, const Tensor &cos_cache, infinicore::nn::RoPE::Algo algo);
-    static common::OpDispatcher<schema> &dispatcher();
-};
 
-// Internal function
-void rope_(Tensor x_out, const Tensor &x, const Tensor &pos, const Tensor &sin_table, const Tensor &cos_table, infinicore::nn::RoPE::Algo algo);
+INFINICORE_GRAPH_OP_CLASS(RoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infinicore::nn::RoPE::Algo);
+
+// Internal
+void rope_(Tensor x_out,
+           const Tensor &x,
+           const Tensor &pos,
+           const Tensor &sin_table,
+           const Tensor &cos_table,
+           infinicore::nn::RoPE::Algo algo);
+
+// Public API
+Tensor rope(const Tensor &x,
+            const Tensor &pos,
+            const Tensor &sin_table,
+            const Tensor &cos_table,
+            infinicore::nn::RoPE::Algo algo);
 
-// Public API that uses infinicore::nn::RoPE::Algo
-Tensor rope(const Tensor &x, const Tensor &pos, const Tensor &sin_table, const Tensor &cos_table, infinicore::nn::RoPE::Algo algo);
 } // namespace infinicore::op

include/infinicore/ops/scaled_mm_i8.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(I8Gemm, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional<Tensor>);
+
+void scaled_mm_i8_(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional<Tensor> bias);
+} // namespace infinicore::op

include/infinicore/ops/silu_and_mul.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(SiluAndMul, Tensor, const Tensor &);
+
+Tensor silu_and_mul(const Tensor &x);
+void silu_and_mul_(Tensor out, const Tensor &x);
+
+} // namespace infinicore::op

include/infinicore/ops/swiglu.hpp

 #pragma once
 
 #include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "../tensor.hpp"
 #include "common/op.hpp"
 
 namespace infinicore::op {
-class SwiGLU {
-public:
-    using schema = void (*)(Tensor, Tensor, Tensor);
-    static void execute(Tensor c, Tensor a, Tensor b);
-    static common::OpDispatcher<schema> &dispatcher();
-};
 
-Tensor swiglu(Tensor a, Tensor b);
-void swiglu_(Tensor c, Tensor a, Tensor b);
+INFINICORE_GRAPH_OP_CLASS(SwiGLU, Tensor, const Tensor &, const Tensor &);
+
+Tensor swiglu(const Tensor &a, const Tensor &b);
+void swiglu_(Tensor c, const Tensor &a, const Tensor &b);
+
 } // namespace infinicore::op

include/infinicore/quantization.hpp

+#pragma once
+
+#include "quantization/awq.hpp"
+#include "quantization/base_quantization.hpp"
+#include "quantization/compressed_tensors.hpp"
+#include "quantization/none_quantizaiton.hpp"
+#include "quantization/quantization_scheme.hpp"

include/infinicore/quantization/awq.hpp

+#pragma once
+#include "base_quantization.hpp"
+namespace infinicore::quantization {
+
+class AWQ : public BaseQuantization {
+    // This is a temporary class that currently only returns AWQ_W4A16.
+    // Future enhancements should parse quant_config to extract detailed quantization
+    // information and support multiple quantization schemes.
+public:
+    explicit AWQ(const nlohmann::json &quant_config)
+        : BaseQuantization(quant_config) {};
+
+    infinicore::quantization::QuantScheme
+    get_quant_scheme() const override {
+        return infinicore::quantization::QuantScheme::AWQ_W4A16;
+    };
+};
+
+} // namespace infinicore::quantization

include/infinicore/quantization/base_quantization.hpp

+#pragma once
+#include "nlohmann/json.hpp"
+#include "quantization_scheme.hpp"
+
+namespace infinicore::quantization {
+class BaseQuantization {
+    // Base class for quantization schemes. Intended to be extended to support various quantization methods.
+public:
+    explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config) {};
+    virtual ~BaseQuantization() = default;
+
+    virtual infinicore::quantization::QuantScheme get_quant_scheme() const = 0;
+
+protected:
+    nlohmann::json quant_config_;
+};
+} // namespace infinicore::quantization

include/infinicore/quantization/compressed_tensors.hpp

+#pragma once
+
+#include "base_quantization.hpp"
+namespace infinicore::quantization {
+
+class CompressedTensors : public BaseQuantization {
+    // This is a temporary class that currently only returns COMPRESSED_TENSOR_W8A8I8.
+    // Future enhancements should parse quant_config to extract detailed quantization
+    // information and support multiple quantization schemes.
+public:
+    explicit CompressedTensors(const nlohmann::json &quant_config)
+        : BaseQuantization(quant_config) {};
+
+    infinicore::quantization::QuantScheme
+    get_quant_scheme() const override {
+        return infinicore::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8;
+    };
+};
+
+} // namespace infinicore::quantization

include/infinicore/quantization/none_quantizaiton.hpp

+#pragma once
+
+#include "base_quantization.hpp"
+namespace infinicore::quantization {
+
+class NoneQuantization : public BaseQuantization {
+    // This is a temporary class that currently only returns COMPRESSED_TENSOR_W8A8I8.
+    // Future enhancements should parse quant_config to extract detailed quantization
+    // information and support multiple quantization schemes.
+public:
+    explicit NoneQuantization(const nlohmann::json &quant_config)
+        : BaseQuantization(quant_config) {};
+
+    infinicore::quantization::QuantScheme
+    get_quant_scheme() const override {
+        return infinicore::quantization::QuantScheme::NONE;
+    };
+};
+
+} // namespace infinicore::quantization

include/infinicore/quantization/quantization_scheme.hpp

+// quant.hpp
+#pragma once
+
+namespace infinicore::quantization {
+
+enum class QuantScheme {
+    NONE,
+    COMPRESSED_TENSOR_W8A8I8,
+    AWQ_W4A16,
+};
+
+} // namespace infinicore::quantization