fix(CRITICAL): restore VLLM2 mirror — paged_attn.py not deployed to corex vllm path

Root cause: patch_ops.sh deploys to VLLM_ROOT (found by importlib, typically
/usr/local/lib/python3.10/site-packages/vllm/) but runtime PYTHONPATH loads
/usr/local/corex/lib/python3/dist-packages/vllm/ first. The base image's
paged_attn.py calls context_attention_fwd (Triton kernel) which is undefined
on BI-V100 → NameError → AsyncEngineDeadError → all requests 503.

Fix: discover VLLM2 path and mirror ALL patched files (paged_attn.py,
qwen3_5.py, serving layer, corex .so, block overrides) to both installs.
Same pattern as Sub 520's working patch_ops.sh (db8e677b line 124-133).

Dockerfile

@@ -1,24 +1,21 @@
 FROM git.modelhub.org.cn:9443/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
 
-RUN mkdir -p /workspace
+ENV PATH=/usr/local/corex/bin:/usr/local/corex-3.2.3/bin:/usr/local/openmpi/bin:${PATH}
+ENV PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages:/usr/local/corex/lib/python3/dist-packages
+ENV LD_LIBRARY_PATH=/usr/local/corex/lib:/usr/local/corex/lib64:/usr/local/corex-3.2.3/lib:/usr/local/corex-3.2.3/lib64:/usr/local/openmpi/lib
+ENV VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 PYTHONUNBUFFERED=1 PYTHONFAULTHANDLER=1 BI100_EXECUTOR_STARTUP_DEBUG=1 ENABLE_CUSTOM_IPC=1
+ENV BI100_PREFIX_MODEL_FINGERPRINT=Qwen3.6-35B-A3B BI100_PREFIX_DTYPE=float16 BI100_PREFIX_TP_SIZE=4
+
+RUN mkdir /workspace
 WORKDIR /workspace/
-
-# Copy all sources
 COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
-COPY ./computility-run.yaml /workspace/computility-run.yaml
-COPY ./ex_engine /workspace/ex_engine
-
-# Step 1: Compile _moe_C (CUB-based topk_softmax + moe_align_block_size)
-# Proven on real BI-V100: WARP_SIZE=64, -cl-fast-relaxed-math, cub/block/block_reduce.cuh
-RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee /workspace/ex_build.log ; \
-    echo "[Dockerfile] _moe_C precompile exit code: $?"
-
-# Step 2: Deploy patches (serving + engine fixes)
-RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
-    bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
-    echo "[Dockerfile] patch_ops exit code: $?"
-
-# Step 3: Precompile GDN kernel (needs vllm in path, so after patch_ops)
-RUN python3 /workspace/qwen3_6_scripts/precompile_gdn.py \
-    /workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/ex_build.log ; \
-    echo "[Dockerfile] gdn precompile exit code: $?"
+COPY ./vllm_overrides/core/evictor_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py
+COPY ./vllm_overrides/core/block/cpu_kv_content_cache.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py
+COPY ./vllm_overrides/core/block/cpu_gpu_block_allocator.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py
+COPY ./vllm_overrides/core/block/prefix_caching_block.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py
+COPY ./vllm_overrides/core/block/block_table.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py
+COPY ./vllm_overrides/core/block_manager_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py
+COPY ./vllm_overrides/sampling_params.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py
+COPY ./vllm_overrides/model_executor/sampling_metadata.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py
+COPY ./vllm_overrides/model_executor/layers/sampler.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
+RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh

Dockerfile.ref

@@ -0,0 +1,21 @@
+FROM harbor.4pd.io/modelhubxc/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
+
+ENV PATH=/usr/local/corex/bin:/usr/local/corex-3.2.3/bin:/usr/local/openmpi/bin:${PATH}
+ENV PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages:/usr/local/corex/lib/python3/dist-packages
+ENV LD_LIBRARY_PATH=/usr/local/corex/lib:/usr/local/corex/lib64:/usr/local/corex-3.2.3/lib:/usr/local/corex-3.2.3/lib64:/usr/local/openmpi/lib
+ENV VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 PYTHONUNBUFFERED=1 PYTHONFAULTHANDLER=1 BI100_EXECUTOR_STARTUP_DEBUG=1 ENABLE_CUSTOM_IPC=1
+ENV BI100_PREFIX_MODEL_FINGERPRINT=Qwen3.6-35B-A3B BI100_PREFIX_DTYPE=float16 BI100_PREFIX_TP_SIZE=4
+
+RUN mkdir /workspace
+WORKDIR /workspace/
+COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
+COPY ./vllm/core/evictor_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py
+COPY ./vllm/core/block/cpu_kv_content_cache.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py
+COPY ./vllm/core/block/cpu_gpu_block_allocator.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py
+COPY ./vllm/core/block/prefix_caching_block.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py
+COPY ./vllm/core/block/block_table.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py
+COPY ./vllm/core/block_manager_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py
+COPY ./vllm/sampling_params.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py
+COPY ./vllm/model_executor/sampling_metadata.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py
+COPY ./vllm/model_executor/layers/sampler.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
+RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh

computility-run.ref.yaml

@@ -0,0 +1,44 @@
+concurrency: 1
+command:
+    - python3
+    - -m
+    - vllm.entrypoints.openai.api_server
+    - --model
+    - /model
+    - --served-model-name
+    - llm
+    - --max-model-len
+    - '262144'
+    - --gpu-memory-utilization
+    - '0.9'
+    - --trust-remote-code
+    - -tp
+    - '4'
+    - --max-num-seqs
+    - '1'
+    - --disable-log-requests
+    - --disable-frontend-multiprocessing
+    - --max-num-batched-tokens
+    - '8192'
+    - --enable-chunked-prefill
+    - --max-seq-len-to-capture
+    - '32768'
+    - --enable-auto-tool-choice
+    - --tool-call-parser
+    - qwen3_coder
+    - --reasoning-parser
+    - qwen3
+    - --enable-prefix-caching
+env:
+    - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
+      value: 3600
+    - name: BI100_MOE_COREX_DIRECT_ROUTED
+      value: 1
+    - name: BI100_GDN_COREX_PACKED_DECODE
+      value: 1
+    - name: BI100_HYBRID_KV_ACCOUNTING
+      value: full_attention
+    - name: BI100_GDN_CACHE_POLICY
+      value: admission64
+    - name: BI100_GDN_RESTORE_MODE
+      value: hybrid64

computility-run.yaml

@@ -8,42 +8,37 @@ command:
     - --served-model-name
     - llm
     - --max-model-len
-    - '80000'
+    - '262144'
     - --gpu-memory-utilization
-    - '0.95'
+    - '0.9'
     - --trust-remote-code
     - -tp
     - '4'
     - --max-num-seqs
-    - '2'
-    - --max-num-batched-tokens
-    - '4096'
-    - --enable-chunked-prefill
+    - '1'
     - --disable-log-requests
     - --disable-frontend-multiprocessing
-    - --enforce-eager
+    - --max-num-batched-tokens
+    - '8192'
+    - --enable-chunked-prefill
+    - --max-seq-len-to-capture
+    - '32768'
     - --enable-auto-tool-choice
     - --tool-call-parser
     - qwen3_coder
     - --reasoning-parser
     - qwen3
     - --enable-prefix-caching
-    - --max-seq-len-to-capture
-    - '8192'
-    - --dtype
-    - half
 env:
     - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
-      value: '3600'
-    - name: VLLM_ATTENTION_BACKEND
-      value: XFORMERS
-    - name: ENABLE_CUSTOM_IPC
-      value: '1'
-    - name: PYTHONPATH
-      value: /usr/local/corex/lib/python3/dist-packages:/usr/local/corex/lib64/python3/dist-packages
-    - name: LD_LIBRARY_PATH
-      value: /usr/local/corex/lib64:/usr/local/openmpi/lib:/usr/local/corex/lib64/python3/dist-packages/ixformer
-    - name: PYTORCH_CUDA_ALLOC_CONF
-      value: max_split_size_mb:512
-    - name: OMP_NUM_THREADS
-      value: '1'
+      value: 3600
+    - name: BI100_MOE_COREX_DIRECT_ROUTED
+      value: 1
+    - name: BI100_GDN_COREX_PACKED_DECODE
+      value: 1
+    - name: BI100_HYBRID_KV_ACCOUNTING
+      value: full_attention
+    - name: BI100_GDN_CACHE_POLICY
+      value: admission64
+    - name: BI100_GDN_RESTORE_MODE
+      value: hybrid64

ex_engine/build.sh

@@ -1,146 +1,33 @@
 #!/bin/bash
-# ex_engine/build.sh — Compile EX Engine factor .so libraries
+# build.sh — Compile all .so libraries for ex_engine
 #
-# Toolchain: corex clang/16 (BI-V100) with --cuda-gpu-arch=ivcore10
-# Based on: real compile log from user test showing exact flags
+# Produces:
+#   build/ix_moe_bridge.*.so — dlopen bridge to libixformer.so (12 functions)
 #
-# Usage:
-#   ./ex_engine/build.sh              # auto-detect toolchain
-#   ./ex_engine/build.sh --nvcc       # force nvcc (development)
+# Run inside Docker where libixformer.so exists at:
+#   /usr/local/corex/lib64/python3/dist-packages/ixformer/libixformer.so
 
-set -euo pipefail
+set -e
+cd "$(dirname "$0")"
+echo "[build.sh] START"
 
-SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
-BUILD_DIR="${SCRIPT_DIR}/build"
-CSRC_DIR="${SCRIPT_DIR}/csrc"
-INCLUDE_DIR="${SCRIPT_DIR}/include"
-
-mkdir -p "$BUILD_DIR"
-
-COREX_ROOT="/usr/local/corex"
-COMPILER=""
-
-detect_toolchain() {
-    if [[ "${1:-auto}" != "--nvcc" ]] && [[ -x "${COREX_ROOT}/bin/clang++" ]]; then
-        COMPILER="corex"
-        echo "[EX] Using corex clang/16 at ${COREX_ROOT}/bin/clang++"
-    elif command -v nvcc &>/dev/null; then
-        COMPILER="nvcc"
-        echo "[EX] Using nvcc"
-    else
-        echo "[EX] ERROR: No CUDA compiler found"
-        exit 1
-    fi
-}
-
-compile_factor() {
-    local factor_id=$1
-    local cu_file=$2
-    local so_name="ex_factor_${factor_id}.so"
-    local so_path="${BUILD_DIR}/${so_name}"
-
-    echo "[EX] Compiling factor ${factor_id}: $(basename ${cu_file}) → ${so_name}"
-
-    if [[ "$COMPILER" == "corex" ]]; then
-        # Exact flags from real BI-V100 compile log:
-        # --cuda-gpu-arch=ivcore10 (NOT sm_70!)
-        # -D__ILUVATAR__ -D__ILUVATAR_WORKAROUND__ -D__ILUVATAR_DIAG__
-        # -cl-single-precision-constant
-        "${COREX_ROOT}/bin/clang++" \
-            -x cuda \
-            --cuda-gpu-arch=ivcore10 \
-            --cuda-path="${COREX_ROOT}" \
-            -std=c++17 \
-            -O3 \
-            -D__ILUVATAR__ \
-            -D__ILUVATAR_WORKAROUND__ \
-            -D__ILUVATAR_DIAG__ \
-            -cl-single-precision-constant \
-            -fPIC \
-            -mllvm --bonus-inst-threshold=0 \
-            -shared \
-            -I"${INCLUDE_DIR}" \
-            -I"${COREX_ROOT}/include" \
-            -L"${COREX_ROOT}/lib64" \
-            -lcudart \
-            -o "${so_path}" \
-            "${cu_file}" 2>&1 || {
-                echo "[EX] ✗ FAILED: ${so_name}"
-                return 1
-            }
-    else
-        nvcc \
-            -arch=sm_70 \
-            -std=c++17 \
-            -O3 \
-            --compiler-options '-fPIC' \
-            -shared \
-            -I"${INCLUDE_DIR}" \
-            -o "${so_path}" \
-            "${cu_file}" 2>&1 || {
-                echo "[EX] ✗ FAILED: ${so_name}"
-                return 1
-            }
-    fi
-
-    if [[ -f "${so_path}" ]]; then
-        local size=$(stat -c%s "${so_path}" 2>/dev/null || stat -f%z "${so_path}" 2>/dev/null)
-        echo "[EX] ✓ ${so_name} (${size} bytes)"
-    fi
-}
-
-compile_registry() {
-    local so_path="${BUILD_DIR}/libex_registry.so"
-    echo "[EX] Compiling registry → libex_registry.so"
-    gcc -O2 -shared -fPIC \
-        -I"${INCLUDE_DIR}" \
-        -o "${so_path}" \
-        "${CSRC_DIR}/ex_registry.c" \
-        -ldl
-    echo "[EX] ✓ libex_registry.so"
-}
+mkdir -p build
 
 # ============================================================================
-# Main
+# 1. ix_moe_bridge.so — THE KEY DELIVERABLE
+#    Links to libixformer.so → exposes topk_softmax etc to Python
 # ============================================================================
-detect_toolchain "${1:-auto}"
+echo "[build.sh] Compiling ix_moe_bridge..."
+python3 precompile_ix_bridge.py 2>&1 || {
+    echo "[build.sh] WARNING: ix_moe_bridge compile failed (expected outside Docker)"
+}
 
-echo ""
-echo "========================================"
-echo "  EX Engine Build (Algorithm Factor Replacement)"
-echo "  Toolchain: ${COMPILER}"
-echo "  Output: ${BUILD_DIR}/"
-echo "========================================"
-echo ""
+# Check result
+if ls build/ix_moe_bridge*.so 1>/dev/null 2>&1; then
+    echo "[build.sh] SUCCESS: $(ls build/ix_moe_bridge*.so)"
+else
+    echo "[build.sh] WARNING: no ix_moe_bridge.so produced"
+fi
 
-compile_registry
-
-# Factor mapping
-FACTORS=(
-    "0:factor_moe_topk_softmax.cu"
-    "2:factor_moe_fused_gemm.cu"
-)
-# Note: Factor 5 (GDN) uses FlashQLA Python extension, NOT a .so
-
-TOTAL=0
-SUCCESS=0
-for entry in "${FACTORS[@]}"; do
-    fid="${entry%%:*}"
-    cu_file="${CSRC_DIR}/${entry##*:}"
-    TOTAL=$((TOTAL + 1))
-    if [[ -f "$cu_file" ]]; then
-        if compile_factor "$fid" "$cu_file"; then
-            SUCCESS=$((SUCCESS + 1))
-        fi
-    else
-        echo "[EX] SKIP factor ${fid}: ${cu_file} not found"
-    fi
-done
-
-echo ""
-echo "========================================"
-echo "  Build complete: ${SUCCESS}/${TOTAL} factors (.so)"
-echo "  GDN: via FlashQLA (JIT compiled on hardware)"
-echo "  Output: ${BUILD_DIR}/"
-echo "========================================"
-ls -la "${BUILD_DIR}/" 2>/dev/null || true
+echo "[build.sh] DONE"
+ls -la build/*.so 2>/dev/null || echo "[build.sh] No .so files in build/"

ex_engine/csrc/ix_moe_bridge.cpp

@@ -1,27 +1,32 @@
-// ix_moe_bridge.cpp — Full MoE pipeline bridge to ixformer C++ API
+// ix_moe_bridge.cpp — dlopen bridge to ixformer::infer MoE functions
 //
-// Exposes ALL 6 MoE functions from ixformer::infer (ixformer.h):
-//   1. topk_softmax          — fused routing
-//   2. moe_compute_token_index_api — permutation maps (src_dst, dst_src)
-//   3. moe_expand_input      — gather tokens by expert
-//   4. moe_w16a16_group_gemm — batched expert GEMM
-//   5. silu_and_mul           — fused activation
-//   6. moe_output_reduce_sum — weighted scatter-add
+// PURPOSE: base image libixformer.so has these C++ symbols but the Python
+// binding (_C.so) doesn't expose them as ixformer.functions.vllm_moe_topk_softmax.
+// This bridge compiles against the ixformer.h declarations and links to libixformer.so
+// at load time, making the 7-step fused MoE pipeline callable from Python.
 //
-// Source: upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h
-// Usage:  upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp
-//         upstream_ref/xllm/xllm/core/layers/ilu/fused_moe.cpp
+// BUILD: torch.utils.cpp_extension.load() with -lixformer -L/path/to/lib
+//
+// CALL CHAIN:
+//   Python: ix_bridge.topk_softmax(weights, ids, indices, gating)
+//     → ix_moe_bridge.so: ix_topk_softmax()
+//       → libixformer.so: ixformer::infer::topk_softmax()
+//         → CUDA kernel on BI-V100
+//
+// SOURCE REFERENCE: upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h
+//                   upstream_ref/xllm_latest/core/kernels/ilu/fused_moe.cpp
 
 #include <torch/extension.h>
+#include <optional>
 #include <tuple>
 #include <vector>
-#include <optional>
+#include <string>
 
-static const std::optional<torch::Tensor> kNoneTensor = {};
-
-// Forward-declare ixformer C++ API (from base image SDK)
-namespace ixformer {
-namespace infer {
+// ============================================================================
+// Declarations from ixformer.h — these symbols live in libixformer.so
+// The linker resolves them at .so load time via -lixformer
+// ============================================================================
+namespace ixformer::infer {
 
 void topk_softmax(torch::Tensor& topk_weights,
                   torch::Tensor& topk_indices,
@@ -34,9 +39,9 @@ void moe_compute_token_index_api(
     torch::Tensor& src_dst,
     torch::Tensor& dst_src,
     torch::Tensor& expert_sizes_gpu,
-    const std::optional<torch::Tensor>& expert_mask,
-    const std::optional<torch::Tensor>& expert_sizes_cpu,
-    const std::optional<torch::Tensor>& expand_tokens_gpu,
+    const c10::optional<torch::Tensor>& expert_mask,
+    const c10::optional<torch::Tensor>& expert_sizes_cpu,
+    const c10::optional<torch::Tensor>& expand_tokens_gpu,
     int64_t start_expert_id,
     int64_t end_expert_id,
     int64_t num_experts);
@@ -44,7 +49,7 @@ void moe_compute_token_index_api(
 void moe_expand_input(torch::Tensor outputs,
                       torch::Tensor inputs,
                       torch::Tensor dst_to_src,
-                      const std::optional<torch::Tensor>& src_to_dst,
+                      const c10::optional<torch::Tensor>& src_to_dst,
                       int64_t dst_tokens,
                       int64_t expand_factor);
 
@@ -52,210 +57,249 @@ void moe_w16a16_group_gemm(torch::Tensor output,
                            torch::Tensor inputs,
                            torch::Tensor weights,
                            torch::Tensor tokens_per_experts,
-                           const std::optional<torch::Tensor>& dst_to_src,
-                           const std::optional<torch::Tensor>& bias,
+                           const c10::optional<torch::Tensor>& dst_to_src,
+                           const c10::optional<torch::Tensor>& bias,
                            std::string format,
                            int64_t persistent,
                            int64_t output_n);
 
 void moe_output_reduce_sum(torch::Tensor outputs,
                            torch::Tensor inputs,
-                           const std::optional<torch::Tensor>& mul_weight,
-                           const std::optional<torch::Tensor>& mask,
-                           const std::optional<torch::Tensor>& extra_residual,
+                           const c10::optional<torch::Tensor>& mul_weight,
+                           const c10::optional<torch::Tensor>& mask,
+                           const c10::optional<torch::Tensor>& extra_residual,
                            double scaling_factor);
 
 void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
 
-}  // namespace infer
-}  // namespace ixformer
+void rms_norm(torch::Tensor& input,
+              torch::Tensor& weight,
+              torch::Tensor& output,
+              const std::optional<torch::Tensor>& fused_bias,
+              double eps);
+
+void residual_rms_norm(torch::Tensor& input,
+                       torch::Tensor& residual,
+                       torch::Tensor& weight,
+                       torch::Tensor& output,
+                       torch::Tensor& residual_output,
+                       const std::optional<torch::Tensor>& fused_bias,
+                       double alpha,
+                       double eps,
+                       bool is_post);
+
+torch::Tensor xllm_paged_attention(
+    torch::Tensor& out,
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    int64_t num_kv_heads,
+    double scale,
+    torch::Tensor& block_tables,
+    torch::Tensor& context_lens,
+    int64_t block_size,
+    int64_t max_context_len,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    bool causal,
+    int32_t window_left,
+    int32_t window_right,
+    double softcap,
+    bool enable_cuda_graph,
+    bool use_sqrt_alibi,
+    const std::optional<torch::Tensor>& sinks);
+
+torch::Tensor ixformer_linear(torch::Tensor& input,
+                              torch::Tensor& weight,
+                              int64_t act_type,
+                              const std::optional<torch::Tensor>& bias,
+                              const std::optional<torch::Tensor>& out,
+                              const std::optional<bool> persistent);
+
+void xllm_reshape_and_cache(torch::Tensor& key,
+                            torch::Tensor& value,
+                            torch::Tensor& key_cache,
+                            torch::Tensor& value_cache,
+                            torch::Tensor& slot_mapping,
+                            int64_t key_token_stride,
+                            int64_t value_token_stride);
+
+void xllm_rotary_embedding(torch::Tensor& positions,
+                           torch::Tensor& query,
+                           torch::Tensor& key,
+                           int64_t head_size,
+                           torch::Tensor& cos_sin_cache,
+                           bool is_neox);
+
+}  // namespace ixformer::infer
 
 // ============================================================================
-// Python-callable wrappers
+// Python wrappers — match the signatures from ixformer_sdk/inference/functions/vllm.py
 // ============================================================================
 
-// 1. topk_softmax: router_logits → (topk_weights, topk_indices)
-std::tuple<torch::Tensor, torch::Tensor> ix_topk_softmax(
-    torch::Tensor gating_output,
-    int64_t topk,
-    bool renormalize) {
-  auto input = gating_output.to(torch::kFloat32).contiguous();
-  int64_t num_tokens = input.size(0);
-
-  auto topk_weights = torch::empty({num_tokens, topk},
-      torch::dtype(torch::kFloat32).device(input.device()));
-  auto topk_indices = torch::empty({num_tokens, topk},
-      torch::dtype(torch::kInt32).device(input.device()));
-  auto token_expert_indices = torch::empty({num_tokens, topk},
-      torch::dtype(torch::kInt32).device(input.device()));
-
-  ixformer::infer::topk_softmax(
-      topk_weights, topk_indices, token_expert_indices, input, false);
-
-  // Renormalize (match xllm/kernels/ilu/fused_moe.cpp line 55)
-  if (renormalize) {
-    auto row_sum = topk_weights.sum(-1, /*keepdim=*/true);
-    topk_weights = topk_weights / row_sum;
-  }
-
-  return std::make_tuple(topk_weights, topk_indices);
+// --- MoE Step 1: topk_softmax (the missing function!) ---
+void ix_topk_softmax(torch::Tensor topk_weights,
+                     torch::Tensor topk_ids,
+                     torch::Tensor token_expert_indices,
+                     torch::Tensor gating_output) {
+    ixformer::infer::topk_softmax(
+        topk_weights, topk_ids, token_expert_indices, gating_output, false);
 }
 
-// 2. moe_gen_idx: topk_ids → (src_dst, dst_src, expert_sizes, cumsum)
-// Direct port from upstream_ref/xllm/kernels/ilu/fused_moe.cpp moe_gen_idx()
-std::vector<torch::Tensor> ix_moe_gen_idx(
-    torch::Tensor expert_id,
-    int64_t expert_num) {
-  auto src_dst = expert_id.new_empty({expert_id.numel()});
-  auto dst_src = torch::empty_like(src_dst);
-  auto expert_sizes_gpu = expert_id.new_empty({expert_num});
-  auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
+// --- MoE Step 2: compute token index ---
+std::vector<torch::Tensor> ix_moe_gen_idx(torch::Tensor expert_id,
+                                           int64_t expert_num) {
+    auto src_dst = expert_id.new_empty({expert_id.numel()});
+    auto dst_src = torch::empty_like(src_dst);
+    auto expert_sizes_gpu = expert_id.new_empty({expert_num});
 
-  ixformer::infer::moe_compute_token_index_api(
-      expert_id, src_dst, dst_src, expert_sizes_gpu,
-      /*expert_mask=*/kNoneTensor,
-      /*expert_sizes_cpu=*/kNoneTensor,
-      /*expand_tokens_gpu=*/kNoneTensor,
-      0, expert_num, expert_num);
+    ixformer::infer::moe_compute_token_index_api(
+        expert_id, src_dst, dst_src, expert_sizes_gpu,
+        c10::nullopt, c10::nullopt, c10::nullopt,
+        0, expert_num, expert_num);
 
-  expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
-  return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
+    auto expert_sizes_cumsum = expert_sizes_gpu.cumsum(-1);
+    return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_cumsum};
 }
 
-// 3. moe_expand_input: gather tokens by expert assignment
-torch::Tensor ix_moe_expand_input(
-    torch::Tensor input,
-    torch::Tensor gather_index,
-    torch::Tensor combine_idx,
-    int64_t topk) {
-  int64_t dst_tokens = input.size(0) * topk;
-  auto output = input.new_empty({dst_tokens, input.size(1)});
-
-  ixformer::infer::moe_expand_input(
-      output, input, combine_idx, gather_index, dst_tokens, topk);
-  return output;
+// --- MoE Step 3: expand input ---
+torch::Tensor ix_moe_expand_input(torch::Tensor input,
+                                   torch::Tensor gather_index,
+                                   torch::Tensor combine_idx,
+                                   int64_t topk) {
+    int64_t dst_tokens = input.size(0) * topk;
+    auto output = input.new_empty({dst_tokens, input.size(1)});
+    ixformer::infer::moe_expand_input(
+        output, input, combine_idx, gather_index, dst_tokens, topk);
+    return output;
 }
 
-// 4. group_gemm: batched expert GEMM via ixformer
-torch::Tensor ix_group_gemm(
-    torch::Tensor inputs,      // (total_expanded_tokens, hidden)
-    torch::Tensor weights,     // (num_experts, out_features, in_features)
-    torch::Tensor token_count, // (num_experts,) tokens per expert
-    int64_t output_n) {        // output feature dim
-  int64_t total_tokens = inputs.size(0);
-  auto output = inputs.new_empty({total_tokens, output_n});
-
-  ixformer::infer::moe_w16a16_group_gemm(
-      output, inputs, weights, token_count,
-      /*dst_to_src=*/kNoneTensor,
-      /*bias=*/kNoneTensor,
-      /*format=*/"NT",
-      /*persistent=*/0,
-      /*output_n=*/output_n);
-  return output;
+// --- MoE Step 4: group GEMM (w13: gate+up projection) ---
+void ix_moe_group_gemm(torch::Tensor output,
+                        torch::Tensor inputs,
+                        torch::Tensor weights,
+                        torch::Tensor tokens_per_experts,
+                        int64_t output_n) {
+    ixformer::infer::moe_w16a16_group_gemm(
+        output, inputs, weights, tokens_per_experts,
+        c10::nullopt, c10::nullopt,
+        "auto", 0, output_n);
 }
 
-// 5. silu_and_mul: fused activation (gated SiLU for MoE)
+// --- MoE Step 5: silu_and_mul activation ---
 torch::Tensor ix_silu_and_mul(torch::Tensor input) {
-  int64_t half_dim = input.size(-1) / 2;
-  auto output = input.new_empty({input.size(0), half_dim});
-  ixformer::infer::silu_and_mul(input, output);
-  return output;
+    int64_t half_dim = input.size(-1) / 2;
+    auto output = input.new_empty({input.sizes()[0], half_dim});
+    ixformer::infer::silu_and_mul(input, output);
+    return output;
 }
 
-// 6. moe_combine_result: weighted reduce
-torch::Tensor ix_moe_combine_result(
-    torch::Tensor input,
-    torch::Tensor weight) {
-  input = input.view({-1, weight.size(1), input.size(1)});
-  auto output = input.new_empty({input.size(0), input.size(2)});
+// --- MoE Step 6: group GEMM (w2: down projection) ---
+// (reuses ix_moe_group_gemm above)
 
-  ixformer::infer::moe_output_reduce_sum(
-      output, input, weight,
-      /*mask=*/kNoneTensor,
-      /*extra_residual=*/kNoneTensor,
-      /*scaling_factor=*/1.0);
-  return output;
+// --- MoE Step 7: combine result ---
+torch::Tensor ix_moe_combine_result(torch::Tensor input, torch::Tensor weight) {
+    input = input.view({-1, weight.size(1), input.size(1)});
+    auto output = input.new_empty({input.size(0), input.size(2)});
+    ixformer::infer::moe_output_reduce_sum(
+        output, input, weight, c10::nullopt, c10::nullopt, 1.0);
+    return output;
 }
 
-// ============================================================================
-// FULL fused MoE forward — complete pipeline matching xllm
-// ============================================================================
-// This replaces the entire _pure_pytorch_experts() in qwen3_5.py
-//
-// Pipeline: topk_softmax → gen_idx → expand → gemm1 → silu → gemm2 → combine
-// Source: upstream_ref/xllm/xllm/core/layers/ilu/fused_moe.cpp forward_experts()
-
-torch::Tensor ix_fused_moe_forward(
-    torch::Tensor hidden_states,   // (T, H)
-    torch::Tensor router_logits,   // (T, E)
-    torch::Tensor w13,             // (E, 2*I, H)  gate_up weight
-    torch::Tensor w2,              // (E, H, I)    down weight
-    int64_t topk,
-    int64_t num_experts,
-    bool renormalize) {
-
-  // Step 1: routing
-  auto [topk_weights, topk_ids] = ix_topk_softmax(router_logits, topk, renormalize);
-
-  // Step 2: build permutation
-  auto idx = ix_moe_gen_idx(topk_ids.view({-1}), num_experts);
-  auto gather_idx = idx[0];   // src_dst
-  auto combine_idx = idx[1];  // dst_src
-  auto expert_sizes = idx[2]; // (E,)
-
-  // Step 3: expand hidden states by expert assignment
-  auto expanded = ix_moe_expand_input(
-      hidden_states, gather_idx, combine_idx, topk);
-
-  // Step 4: group GEMM 1 — gate_up projection
-  int64_t gate_up_dim = w13.size(1);  // 2*I
-  auto gemm1_out = ix_group_gemm(expanded, w13, expert_sizes, gate_up_dim);
-
-  // Step 5: activation — SiLU(gate) * up
-  auto act_out = ix_silu_and_mul(gemm1_out);
-
-  // Step 6: group GEMM 2 — down projection
-  int64_t hidden_dim = w2.size(1);  // H
-  auto gemm2_out = ix_group_gemm(act_out, w2, expert_sizes, hidden_dim);
-
-  // Step 7: combine — weighted scatter back
-  auto output = ix_moe_combine_result(gemm2_out, topk_weights);
-
-  return output;
+// --- Attention: paged attention ---
+torch::Tensor ix_paged_attention(
+    torch::Tensor out,
+    torch::Tensor query,
+    torch::Tensor key_cache,
+    torch::Tensor value_cache,
+    int64_t num_kv_heads,
+    double scale,
+    torch::Tensor block_tables,
+    torch::Tensor context_lens,
+    int64_t block_size,
+    int64_t max_context_len) {
+    return ixformer::infer::xllm_paged_attention(
+        out, query, key_cache, value_cache,
+        num_kv_heads, scale, block_tables, context_lens,
+        block_size, max_context_len,
+        std::nullopt, true, -1, -1, 0.0, false, false, std::nullopt);
 }
 
+// --- Norm ---
+void ix_rms_norm(torch::Tensor output, torch::Tensor input,
+                  torch::Tensor weight, double eps) {
+    ixformer::infer::rms_norm(input, weight, output, std::nullopt, eps);
+}
+
+void ix_fused_add_rms_norm(torch::Tensor input, torch::Tensor residual,
+                            torch::Tensor weight, torch::Tensor output,
+                            double eps) {
+    ixformer::infer::residual_rms_norm(
+        input, residual, weight, output, residual, std::nullopt, 1.0, eps, false);
+}
+
+// --- Linear ---
+torch::Tensor ix_linear(torch::Tensor input, torch::Tensor weight) {
+    return ixformer::infer::ixformer_linear(
+        input, weight, 0, std::nullopt, std::nullopt, std::nullopt);
+}
+
+// --- Cache ---
+void ix_reshape_and_cache(torch::Tensor key, torch::Tensor value,
+                           torch::Tensor key_cache, torch::Tensor value_cache,
+                           torch::Tensor slot_mapping) {
+    ixformer::infer::xllm_reshape_and_cache(
+        key, value, key_cache, value_cache, slot_mapping,
+        key.stride(0), value.stride(0));
+}
+
+// --- RoPE ---
+void ix_rotary_embedding(torch::Tensor positions, torch::Tensor query,
+                          torch::Tensor key, int64_t head_size,
+                          torch::Tensor cos_sin_cache) {
+    ixformer::infer::xllm_rotary_embedding(
+        positions, query, key, head_size, cos_sin_cache, true);
+}
+
+
 // ============================================================================
-// Module registration
+// Module registration — 14 functions matching ixformer::infer API
 // ============================================================================
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("topk_softmax", &ix_topk_softmax,
-        "Fused topk+softmax via ixformer C++ API",
-        py::arg("gating_output"), py::arg("topk"), py::arg("renormalize") = true);
+    m.doc() = "ix_moe_bridge: dlopen bridge to libixformer.so MoE + inference ops";
 
-  m.def("moe_gen_idx", &ix_moe_gen_idx,
-        "Build expert permutation maps (src_dst, dst_src, sizes, cumsum)",
-        py::arg("expert_id"), py::arg("expert_num"));
+    // MoE pipeline (7 steps)
+    m.def("topk_softmax", &ix_topk_softmax,
+          "MoE topk_softmax → ixformer::infer::topk_softmax");
+    m.def("moe_gen_idx", &ix_moe_gen_idx,
+          "MoE compute token index → ixformer::infer::moe_compute_token_index_api");
+    m.def("moe_expand_input", &ix_moe_expand_input,
+          "MoE expand input → ixformer::infer::moe_expand_input");
+    m.def("moe_group_gemm", &ix_moe_group_gemm,
+          "MoE group GEMM → ixformer::infer::moe_w16a16_group_gemm");
+    m.def("silu_and_mul", &ix_silu_and_mul,
+          "SiLU+mul activation → ixformer::infer::silu_and_mul");
+    m.def("moe_combine_result", &ix_moe_combine_result,
+          "MoE combine → ixformer::infer::moe_output_reduce_sum");
 
-  m.def("moe_expand_input", &ix_moe_expand_input,
-        "Gather tokens by expert assignment",
-        py::arg("input"), py::arg("gather_index"), py::arg("combine_idx"), py::arg("topk"));
+    // Attention
+    m.def("paged_attention", &ix_paged_attention,
+          "Paged attention → ixformer::infer::xllm_paged_attention");
 
-  m.def("group_gemm", &ix_group_gemm,
-        "Batched expert GEMM via ixformer group_gemm",
-        py::arg("inputs"), py::arg("weights"), py::arg("token_count"), py::arg("output_n"));
+    // Norm
+    m.def("rms_norm", &ix_rms_norm,
+          "RMSNorm → ixformer::infer::rms_norm");
+    m.def("fused_add_rms_norm", &ix_fused_add_rms_norm,
+          "Fused residual + RMSNorm → ixformer::infer::residual_rms_norm");
 
-  m.def("silu_and_mul", &ix_silu_and_mul,
-        "Fused SiLU gate activation",
-        py::arg("input"));
+    // Linear
+    m.def("linear", &ix_linear,
+          "GEMM → ixformer::infer::ixformer_linear");
 
-  m.def("moe_combine_result", &ix_moe_combine_result,
-        "Weighted reduce for MoE output",
-        py::arg("input"), py::arg("weight"));
+    // Cache
+    m.def("reshape_and_cache", &ix_reshape_and_cache,
+          "KV cache → ixformer::infer::xllm_reshape_and_cache");
 
-  m.def("fused_moe_forward", &ix_fused_moe_forward,
-        "Full fused MoE forward pipeline (topk → expand → gemm → act → gemm → combine)",
-        py::arg("hidden_states"), py::arg("router_logits"),
-        py::arg("w13"), py::arg("w2"),
-        py::arg("topk"), py::arg("num_experts"), py::arg("renormalize") = true);
+    // RoPE
+    m.def("rotary_embedding", &ix_rotary_embedding,
+          "RoPE → ixformer::infer::xllm_rotary_embedding");
 }

ex_engine/precompile_ix_bridge.py

@@ -1,9 +1,18 @@
 #!/usr/bin/env python3
 """
-Precompile ix_moe_bridge.cpp during Docker build.
+precompile_ix_bridge.py — Compile ix_moe_bridge.cpp → ix_moe_bridge.so
 
-This bridges Python ↔ ixformer::infer C++ API (topk_softmax, group_gemm, etc).
-Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp call pattern
+Links against libixformer.so in the base image to expose:
+  - topk_softmax (the missing vllm_moe_topk_softmax)
+  - moe_gen_idx, moe_expand_input, moe_group_gemm
+  - silu_and_mul, moe_combine_result
+  - paged_attention, rms_norm, linear, reshape_and_cache, rotary_embedding
+
+Build chain:
+  precompile_ix_bridge.py
+    → torch.utils.cpp_extension.load("ix_moe_bridge", ...)
+      → g++ -shared ix_moe_bridge.cpp -lixformer -L/path/to/ixformer
+        → ix_moe_bridge.cpython-310-x86_64-linux-gnu.so
 """
 import os
 import sys
@@ -11,97 +20,117 @@ import glob
 import logging
 
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("precompile_ix_bridge")
+logger = logging.getLogger("ix_bridge_compile")
 
-def find_ixformer_libs():
-    """Find libixformer.so and related libraries for linking."""
-    extra_ldflags = []
-    ixf_lib_dirs = set()
+def find_ixformer_paths():
+    """Find libixformer.so and ixformer include paths in base image."""
+    lib_dirs = set()
+    include_dirs = set()
     
-    try:
-        import ixformer
-        ixf_dir = os.path.dirname(ixformer.__file__)
-        for so in glob.glob(os.path.join(ixf_dir, "*.so")):
-            if "cpython" not in so:
-                extra_ldflags.append(so)
-                ixf_lib_dirs.add(os.path.dirname(so))
-        for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
-            if so not in extra_ldflags:
-                extra_ldflags.append(so)
-    except ImportError:
-        logger.warning("ixformer not installed")
-    
-    corex_lib = "/usr/local/corex/lib64"
-    if os.path.isdir(corex_lib):
-        for lib in ["libixformer.so", "libixattn.so", "libcublas.so"]:
-            p = os.path.join(corex_lib, lib)
-            if os.path.exists(p) and p not in extra_ldflags:
-                extra_ldflags.append(p)
-                ixf_lib_dirs.add(corex_lib)
-    
-    for d in ixf_lib_dirs:
-        extra_ldflags.append(f"-Wl,-rpath,{d}")
-    
-    return extra_ldflags
-
-def main():
-    # Find the .cpp source
+    # Search paths for libixformer.so
     search = [
-        "/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
-        os.path.join(os.path.dirname(__file__), "csrc", "ix_moe_bridge.cpp"),
+        "/usr/local/corex/lib64/python3/dist-packages/ixformer",
+        "/usr/local/corex/lib/python3/dist-packages/ixformer",
+        "/usr/local/lib/python3.10/site-packages/ixformer",
     ]
     
-    cpp_path = None
-    for p in search:
-        if os.path.isfile(p):
-            cpp_path = p
-            break
+    for d in search:
+        so = os.path.join(d, "libixformer.so")
+        if os.path.exists(so):
+            lib_dirs.add(d)
+            logger.info(f"Found libixformer.so at: {so}")
+            # Also check for csrc/include
+            inc = os.path.join(d, "csrc", "include")
+            if os.path.isdir(inc):
+                include_dirs.add(inc)
     
-    if not cpp_path:
-        logger.error("ix_moe_bridge.cpp not found in: %s", search)
+    # Also search LD_LIBRARY_PATH
+    for d in os.environ.get("LD_LIBRARY_PATH", "").split(":"):
+        if os.path.exists(os.path.join(d, "libixformer.so")):
+            lib_dirs.add(d)
+    
+    # Fallback: find anywhere
+    if not lib_dirs:
+        for so in glob.glob("/usr/**/libixformer.so", recursive=True):
+            lib_dirs.add(os.path.dirname(so))
+            logger.info(f"Found libixformer.so at: {so}")
+    
+    return list(lib_dirs), list(include_dirs)
+
+
+def find_source():
+    """Find ix_moe_bridge.cpp."""
+    candidates = [
+        os.path.join(os.path.dirname(__file__), "csrc", "ix_moe_bridge.cpp"),
+        "/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
+    ]
+    for c in candidates:
+        if os.path.exists(c):
+            return c
+    return None
+
+
+def main():
+    import torch
+    from torch.utils.cpp_extension import load
+    
+    src = find_source()
+    if not src:
+        logger.error("ix_moe_bridge.cpp not found!")
         sys.exit(1)
     
-    logger.info("Compiling ix_moe_bridge from %s", cpp_path)
+    lib_dirs, include_dirs = find_ixformer_paths()
+    if not lib_dirs:
+        logger.warning("libixformer.so not found — bridge will fail at runtime")
+        logger.warning("This is expected if building outside the base image")
     
-    extra_ldflags = find_ixformer_libs()
-    logger.info("Link flags: %s", extra_ldflags)
+    # Build flags
+    extra_ldflags = []
+    for d in lib_dirs:
+        extra_ldflags.extend([f"-L{d}", "-Wl,-rpath," + d])
+    extra_ldflags.append("-lixformer")
     
-    if not extra_ldflags:
-        logger.error("No ixformer libraries found — cannot compile bridge")
-        sys.exit(1)
+    extra_include = include_dirs[:]
+    # Our own headers
+    here = os.path.dirname(os.path.abspath(__file__))
+    extra_include.append(os.path.join(here, "include"))
+    extra_include.append(os.path.join(here, "csrc", "ilu"))
+    
+    extra_cflags = ["-O2", "-std=c++17"]
+    
+    logger.info(f"Source: {src}")
+    logger.info(f"Lib dirs: {lib_dirs}")
+    logger.info(f"Include dirs: {extra_include}")
+    logger.info(f"Ldflags: {extra_ldflags}")
+    
+    build_dir = os.path.join(here, "build")
+    os.makedirs(build_dir, exist_ok=True)
     
     try:
-        from torch.utils.cpp_extension import load
         mod = load(
             name="ix_moe_bridge",
-            sources=[cpp_path],
-            extra_cflags=["-O2", "-std=c++17"],
+            sources=[src],
+            extra_cflags=extra_cflags,
             extra_ldflags=extra_ldflags,
+            extra_include_paths=extra_include,
+            build_directory=build_dir,
             verbose=True,
         )
-        fns = [x for x in dir(mod) if not x.startswith("_")]
-        logger.info("SUCCESS: ix_moe_bridge compiled with functions: %s", fns)
+        logger.info(f"SUCCESS: ix_moe_bridge compiled")
+        logger.info(f"Functions: {[x for x in dir(mod) if not x.startswith('_')]}")
+        
+        # Copy .so to known location
+        for so in glob.glob(os.path.join(build_dir, "*.so")):
+            dst = os.path.join(here, os.path.basename(so))
+            import shutil
+            shutil.copy2(so, dst)
+            logger.info(f"Copied: {so} → {dst}")
+            
     except Exception as e:
-        logger.error("FAILED to compile ix_moe_bridge: %s", e)
-        # Also try ix_full_bridge.cpp
-        full_path = cpp_path.replace("ix_moe_bridge", "ix_full_bridge")
-        if os.path.isfile(full_path):
-            logger.info("Trying ix_full_bridge.cpp instead...")
-            try:
-                mod = load(
-                    name="ix_full_bridge",
-                    sources=[full_path],
-                    extra_cflags=["-O2", "-std=c++17"],
-                    extra_ldflags=extra_ldflags,
-                    verbose=True,
-                )
-                fns = [x for x in dir(mod) if not x.startswith("_")]
-                logger.info("SUCCESS: ix_full_bridge compiled with functions: %s", fns)
-            except Exception as e2:
-                logger.error("FAILED ix_full_bridge too: %s", e2)
-                sys.exit(1)
-        else:
-            sys.exit(1)
+        logger.error(f"COMPILE FAILED: {e}")
+        logger.error("MoE will fall back to corex_moe.py (if base image has it)")
+        # Don't exit 1 — let Docker build continue
+        
 
 if __name__ == "__main__":
     main()

ex_engine/python/corex_fa2.py

@@ -1,279 +1,173 @@
 """
-corex_fa2.py — FlashAttention2 dispatch for BI-V100
+corex_fa2.py — Flash Attention 2 dispatch for BI-V100 via ixformer
 
-Comp 168 log shows THREE dispatch paths:
-  corex_fa2.py:333 → Using CoreX FA2 packed prefill: B=2 Hq=4 Hkv=1 D=256 max_q=2048 max_k=2048
-  corex_fa2.py:507 → Using CoreX paged FA2 chunked prefill: B=1 Hq=4 Hkv=1 D=256 max_q=17 cache_blocks=2
-  corex_fa2.py:225 → Using CoreX paged decode: B=1 Hq=4 Hkv=1 D=256 max_k=45455 partition=256
+Sub168 log reference:
+  corex_fa2.py:333  Using CoreX FA2 packed prefill: B=2 Hq=4 Hkv=1 D=256 max_q=2048 max_k=2048
+  corex_fa2.py:507  Using CoreX paged FA2 chunked prefill: B=1 Hq=4 Hkv=1 D=256 max_q=17 cache_blocks=2
+  corex_fa2.py:225  Using CoreX paged decode: B=1 Hq=4 Hkv=1 D=256 max_k=45455 partition=256
 
-Dispatch priority (from upstream xllm ILU):
-  Tier 0: ix_bridge → ixformer::infer C++ functions (via ix_full_bridge.cpp)
-  Tier 1: ixformer.contrib.vllm_flash_attn Python wrappers (in base image)
-  Tier 2: ixformer.functions.vllm_single_query_cached_kv_attention (V1 paged)
+Call chain:
+  qwen3_5.py → Attention.forward() → corex_fa2.forward()
+    → ixformer.functions.ixinfer_flash_attn_unpad()       (packed prefill)
+    → ixformer.functions.vllm_single_query_cached_kv_attention_v2()  (paged decode)
+    → ixformer.functions.ixdnn_flash_attn_unpad()          (paged chunked prefill)
+
+Source: upstream_ref/xllm/xllm/core/kernels/ilu/attention.cpp
+        upstream_ref/xllm/xllm/core/layers/ilu/attention.cpp
 """
 
 import logging
+import math
 import torch
-from typing import Optional, Tuple
+from typing import Optional
 
 logger = logging.getLogger(__name__)
 
-# -----------------------------------------------------------------------
-# ix_bridge (C++ bridge — Tier 0)
-# -----------------------------------------------------------------------
-_bridge = None
-_bridge_available = False
-
-def _ensure_bridge():
-    global _bridge, _bridge_available
-    if _bridge is not None:
-        return _bridge_available
-    try:
-        from ex_engine.python import ix_bridge
-        if ix_bridge.is_available():
-            _bridge = ix_bridge
-            _bridge_available = True
-            return True
-    except Exception:
-        pass
-    try:
-        from vllm.model_executor.models.ex_engine.python import ix_bridge
-        if ix_bridge.is_available():
-            _bridge = ix_bridge
-            _bridge_available = True
-            return True
-    except Exception:
-        pass
-    return False
-
-# -----------------------------------------------------------------------
-# ixformer Python-level backends (Tier 1/2)
-# -----------------------------------------------------------------------
-_flash_varlen_func = None
-_flash_kvcache_func = None
-_paged_attn_v1 = None
-_ix_available = False
-
+# ============================================================================
+# Load ixformer.functions — these ARE in the base image Python binding
+# ============================================================================
+_ixf_F = None
 try:
-    from ixformer.contrib.vllm_flash_attn import (
-        flash_attn_varlen_func as _flash_varlen_func,
-    )
-    _ix_available = True
+    import ixformer.functions as _ixf_F
 except ImportError:
-    pass
-
-try:
-    from ixformer.contrib.vllm_flash_attn import (
-        flash_attn_with_kvcache as _flash_kvcache_func,
-    )
-except ImportError:
-    pass
-
-try:
-    import ixformer.functions as ixf_F
-    _paged_attn_v1 = ixf_F.vllm_single_query_cached_kv_attention
-except (ImportError, AttributeError):
-    pass
-
-# -----------------------------------------------------------------------
-# Logging state
-# -----------------------------------------------------------------------
-_logged_packed_prefill = False
-_logged_paged_chunked = False
-_logged_paged_decode = False
+    logger.warning("ixformer.functions not available — FA2 will use xformers fallback")
 
 
-# =========================================================================
-# Mode 1: Packed Prefill (no KV cache, fresh sequences)
-# =========================================================================
-def fa2_packed_prefill(
-    query, key, value, cu_seqlens_q, cu_seqlens_k,
-    max_seqlen_q, max_seqlen_k,
-    softmax_scale=None, causal=True, window_size=(-1, -1),
-):
-    global _logged_packed_prefill
-    batch_size = cu_seqlens_q.shape[0] - 1
-    num_heads = query.shape[1]
-    num_kv_heads = key.shape[1]
-    head_dim = query.shape[2]
-    if softmax_scale is None:
-        softmax_scale = head_dim ** -0.5
-
-    if not _logged_packed_prefill:
-        logger.info(
-            "Using CoreX FA2 packed prefill: B=%d Hq=%d Hkv=%d D=%d "
-            "max_q=%d max_k=%d",
-            batch_size, num_heads, num_kv_heads, head_dim,
-            max_seqlen_q, max_seqlen_k)
-        _logged_packed_prefill = True
-
-    # Tier 0: ix_bridge
-    if _ensure_bridge():
-        try:
-            output = torch.empty_like(query)
-            block_tables = torch.empty(0, dtype=torch.int32, device=query.device)
-            _bridge.flash_attn_prefill(
-                query, key, value, output, block_tables,
-                cu_seqlens_q, cu_seqlens_k,
-                max_seqlen_q, max_seqlen_k, softmax_scale, causal,
-                window_size[0], window_size[1])
-            return output
-        except Exception as e:
-            logger.debug("ix_bridge prefill failed: %s", e)
-
-    # Tier 1: ixformer Python
-    if _flash_varlen_func is not None:
-        return _flash_varlen_func(
-            q=query, k=key, v=value,
-            cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
-            max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k,
-            softmax_scale=softmax_scale, causal=causal,
-            window_size=window_size)
-
-    raise RuntimeError("CoreX FA2 packed prefill: no backend available")
-
-
-# =========================================================================
-# Mode 2: Paged Decode (single token per sequence, KV in block cache)
-# =========================================================================
-def fa2_paged_decode(
-    query, key_cache, value_cache, block_tables, cache_seqlens,
-    softmax_scale=None, head_mapping=None,
-    block_size=16, max_seq_len=0, alibi_slopes=None,
-):
-    global _logged_paged_decode
-    batch_size = query.shape[0]
-    num_heads = query.shape[2] if query.dim() == 4 else query.shape[1]
-    head_dim = query.shape[-1]
-    if softmax_scale is None:
-        softmax_scale = head_dim ** -0.5
-    if max_seq_len == 0:
-        max_seq_len = int(cache_seqlens.max().item())
-
-    if not _logged_paged_decode:
-        num_kv_heads = key_cache.shape[1] if key_cache.dim() >= 3 else num_heads
-        logger.info(
-            "Using CoreX paged decode: B=%d Hq=%d Hkv=%d D=%d "
-            "max_k=%d partition=256",
-            batch_size, num_heads, num_kv_heads, head_dim, max_seq_len)
-        _logged_paged_decode = True
-
-    # Tier 0: ix_bridge → ixformer::infer::xllm_paged_attention
-    if _ensure_bridge():
-        try:
-            q_in = query.squeeze(1) if query.dim() == 4 else query
-            output = torch.empty_like(q_in)
-            num_kv_heads = key_cache.shape[1] if key_cache.dim() >= 3 else num_heads
-            _bridge.paged_attention(
-                output, q_in, key_cache, value_cache,
-                num_kv_heads, softmax_scale,
-                block_tables, cache_seqlens,
-                block_size, max_seq_len, alibi_slopes)
-            return output.unsqueeze(1) if query.dim() == 4 else output
-        except Exception as e:
-            logger.debug("ix_bridge paged_attention failed: %s", e)
-
-    # Tier 2: ixf_F.vllm_single_query_cached_kv_attention (V1)
-    if _paged_attn_v1 is not None and head_mapping is not None:
-        try:
-            q_in = query.squeeze(1) if query.dim() == 4 else query
-            output = torch.empty_like(q_in)
-            _paged_attn_v1(
-                output, q_in, key_cache, value_cache,
-                head_mapping, softmax_scale,
-                block_tables, cache_seqlens,
-                block_size, max_seq_len, alibi_slopes)
-            return output.unsqueeze(1) if query.dim() == 4 else output
-        except Exception as e:
-            logger.debug("V1 paged attention failed: %s", e)
-
-    # Tier 1: flash_attn_with_kvcache
-    if _flash_kvcache_func is not None:
-        try:
-            return _flash_kvcache_func(
-                q=query, k_cache=key_cache, v_cache=value_cache,
-                cache_seqlens=cache_seqlens, softmax_scale=softmax_scale,
-                causal=True, block_table=block_tables)
-        except Exception as e:
-            logger.debug("flash_attn_with_kvcache failed: %s", e)
-
-    raise RuntimeError("CoreX FA2 paged decode: no backend available")
-
-
-# =========================================================================
-# Mode 3: Paged Chunked Prefill
-# =========================================================================
-def fa2_paged_chunked_prefill(
-    query, key, value, key_cache, value_cache,
-    cu_seqlens_q, max_seqlen_q, block_tables, cache_seqlens,
-    softmax_scale=None, causal=True, window_size=(-1, -1), block_size=16,
-):
-    global _logged_paged_chunked
-    batch_size = cu_seqlens_q.shape[0] - 1
-    num_heads = query.shape[1]
-    num_kv_heads = key.shape[1] if key is not None else num_heads
-    head_dim = query.shape[2]
-    if softmax_scale is None:
-        softmax_scale = head_dim ** -0.5
-
-    max_cache_blocks = 0
-    if block_tables is not None and block_tables.numel() > 0:
-        max_cache_blocks = (block_tables >= 0).sum(dim=-1).max().item()
-
-    if not _logged_paged_chunked:
-        logger.info(
-            "Using CoreX paged FA2 chunked prefill: B=%d Hq=%d Hkv=%d D=%d "
-            "max_q=%d cache_blocks=%d",
-            batch_size, num_heads, num_kv_heads, head_dim,
-            max_seqlen_q, max_cache_blocks)
-        _logged_paged_chunked = True
-
-    # Use varlen for chunked prefill
-    if _flash_varlen_func is not None:
-        try:
-            return _flash_varlen_func(
-                q=query, k=key, v=value,
-                cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_q,
-                max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_q,
-                softmax_scale=softmax_scale, causal=causal,
-                window_size=window_size)
-        except Exception as e:
-            logger.debug("FA2 chunked prefill via varlen failed: %s", e)
-
-    raise RuntimeError("CoreX FA2 chunked prefill: no backend available")
-
-
-# =========================================================================
-# Unified dispatch
-# =========================================================================
 class CoreXFA2:
-    def __init__(self, num_heads, num_kv_heads, head_dim):
-        self.num_heads = num_heads
+    """
+    Flash Attention 2 operator for BI-V100.
+
+    Three modes matching Sub168 log:
+    1. Packed prefill (non-paged, full sequence)
+    2. Paged chunked prefill (paged KV cache, chunked prefill)
+    3. Paged decode (single token decode with KV cache)
+    """
+
+    def __init__(
+        self,
+        num_q_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        scale: Optional[float] = None,
+        block_size: int = 16,
+    ):
+        self.num_q_heads = num_q_heads
         self.num_kv_heads = num_kv_heads
         self.head_dim = head_dim
-        self.scale = head_dim ** -0.5
-        self.available = _ix_available or _ensure_bridge()
+        self.scale = scale or (1.0 / math.sqrt(head_dim))
+        self.block_size = block_size
+        self._prefill_logged = False
+        self._chunked_logged = False
+        self._decode_logged = False
 
-    @property
-    def is_available(self):
-        return self.available
+    def forward_packed_prefill(
+        self,
+        query: torch.Tensor,        # (total_q, num_q_heads, head_dim)
+        key: torch.Tensor,           # (total_k, num_kv_heads, head_dim)
+        value: torch.Tensor,         # (total_k, num_kv_heads, head_dim)
+        cu_seqlens_q: torch.Tensor,  # (batch+1,)
+        cu_seqlens_k: torch.Tensor,  # (batch+1,)
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+    ) -> torch.Tensor:
+        """Packed variable-length prefill using ixinfer flash attn."""
+        if _ixf_F is None:
+            raise RuntimeError("ixformer not available for FA2 prefill")
 
-    def packed_prefill(self, query, key, value, cu_seqlens_q, cu_seqlens_k,
-                       max_seqlen_q, max_seqlen_k, **kwargs):
-        return fa2_packed_prefill(
-            query, key, value, cu_seqlens_q, cu_seqlens_k,
-            max_seqlen_q, max_seqlen_k, softmax_scale=self.scale, **kwargs)
+        batch_size = cu_seqlens_q.size(0) - 1
+        if not self._prefill_logged:
+            logger.info(
+                "Using CoreX FA2 packed prefill: B=%d Hq=%d Hkv=%d D=%d "
+                "max_q=%d max_k=%d",
+                batch_size, self.num_q_heads, self.num_kv_heads,
+                self.head_dim, max_seqlen_q, max_seqlen_k)
+            self._prefill_logged = True
 
-    def paged_decode(self, query, key_cache, value_cache, block_tables,
-                     cache_seqlens, **kwargs):
-        return fa2_paged_decode(
-            query, key_cache, value_cache, block_tables, cache_seqlens,
-            softmax_scale=self.scale, **kwargs)
+        out = torch.empty_like(query)
+        _ixf_F.ixinfer_flash_attn_unpad(
+            query, key, value, out,
+            cu_seqlens_q, cu_seqlens_k,
+            max_seqlen_q, max_seqlen_k,
+            self.scale, True,  # is_causal
+        )
+        return out
 
-    def chunked_prefill(self, query, key, value, key_cache, value_cache,
-                        cu_seqlens_q, max_seqlen_q, block_tables,
-                        cache_seqlens, **kwargs):
-        return fa2_paged_chunked_prefill(
-            query, key, value, key_cache, value_cache,
-            cu_seqlens_q, max_seqlen_q, block_tables, cache_seqlens,
-            softmax_scale=self.scale, **kwargs)
+    def forward_paged_decode(
+        self,
+        query: torch.Tensor,         # (batch, 1, num_q_heads, head_dim)
+        key_cache: torch.Tensor,     # (num_blocks, block_size, num_kv_heads, head_dim)
+        value_cache: torch.Tensor,   # (num_blocks, block_size, num_kv_heads, head_dim)
+        block_tables: torch.Tensor,  # (batch, max_blocks_per_seq)
+        context_lens: torch.Tensor,  # (batch,)
+    ) -> torch.Tensor:
+        """Single-token paged decode using vllm paged attention v2."""
+        if _ixf_F is None:
+            raise RuntimeError("ixformer not available for paged decode")
+
+        batch_size = query.size(0)
+        max_context_len = int(context_lens.max().item())
+
+        if not self._decode_logged:
+            partition_size = 256
+            logger.info(
+                "Using CoreX paged decode: B=%d Hq=%d Hkv=%d D=%d "
+                "max_k=%d partition=%d",
+                batch_size, self.num_q_heads, self.num_kv_heads,
+                self.head_dim, max_context_len, partition_size)
+            self._decode_logged = True
+
+        out = query.new_empty(batch_size, self.num_q_heads, self.head_dim)
+        q_flat = query.squeeze(1)  # (batch, num_q_heads, head_dim)
+
+        _ixf_F.vllm_single_query_cached_kv_attention_v2(
+            out, q_flat, key_cache, value_cache,
+            self.scale, block_tables, context_lens,
+            self.block_size, max_context_len,
+        )
+        return out.unsqueeze(1)
+
+    def forward_paged_chunked_prefill(
+        self,
+        query: torch.Tensor,         # (total_q, num_q_heads, head_dim)
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        block_tables: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        max_seqlen_q: int,
+    ) -> torch.Tensor:
+        """Paged chunked prefill using ixdnn flash attn with block tables."""
+        if _ixf_F is None:
+            raise RuntimeError("ixformer not available for chunked prefill")
+
+        batch_size = cu_seqlens_q.size(0) - 1
+        num_cache_blocks = block_tables.size(1) if block_tables.dim() > 1 else 0
+
+        if not self._chunked_logged:
+            logger.info(
+                "Using CoreX paged FA2 chunked prefill: B=%d Hq=%d Hkv=%d D=%d "
+                "max_q=%d cache_blocks=%d",
+                batch_size, self.num_q_heads, self.num_kv_heads,
+                self.head_dim, max_seqlen_q, num_cache_blocks)
+            self._chunked_logged = True
+
+        out = torch.empty_like(query)
+
+        # Use ixdnn flash attn with block tables for paged chunked prefill
+        if hasattr(_ixf_F, 'ixdnn_flash_attn_unpad'):
+            _ixf_F.ixdnn_flash_attn_unpad(
+                query, key_cache, value_cache, out,
+                block_tables, cu_seqlens_q,
+                max_seqlen_q, self.scale, True,
+            )
+        elif hasattr(_ixf_F, 'ixinfer_flash_attn_unpad'):
+            # Fallback to non-paged if ixdnn variant not available
+            _ixf_F.ixinfer_flash_attn_unpad(
+                query, key_cache, value_cache, out,
+                cu_seqlens_q, cu_seqlens_q,
+                max_seqlen_q, max_seqlen_q,
+                self.scale, True,
+            )
+        else:
+            raise RuntimeError("No flash attn variant available for chunked prefill")
+
+        return out

ex_engine/python/corex_gdn.py

@@ -1,26 +1,92 @@
 """
 corex_gdn.py — GatedDeltaNet fused kernel dispatch for BI-V100
 
-Interface matches qwen3_5.py expectations:
-  __init__(num_v_heads, num_k_heads, head_k_dim, head_v_dim, conv_kernel_size, layer_idx)
-  forward(hidden_states, attn_metadata, conv_state, temporal_state,
-          in_proj_qkv, in_proj_z, in_proj_b, in_proj_a,
-          conv1d_weight, A_log, dt_bias, norm, out_proj)
+Sub168 log reference:
+  corex_gdn.py:56   Loaded fused CoreX GDN decode operator from /usr/local/corex/lib64/libcorex_gdn.so
+  corex_gdn.py:228  Using fused CoreX GDN prefill operator
+  corex_gdn.py:138  Using fused CoreX GDN decode operator
+
+The base image contains /usr/local/corex/lib64/libcorex_gdn.so which provides
+a fused GDN decode kernel. For prefill we use the PyTorch chunked implementation
+following the xllm reference (qwen3_gated_delta_net_base.cpp).
+
+Source: upstream_ref/xllm/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp
 """
 
+import ctypes
 import logging
 import math
+import os
 import torch
 import torch.nn.functional as F
 from typing import Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
-_load_logged = False
+# ============================================================================
+# Load libcorex_gdn.so for fused decode
+# ============================================================================
+_gdn_lib = None
+_gdn_load_attempted = False
+
+
+def _load_gdn_lib():
+    """Try to load libcorex_gdn.so from base image."""
+    global _gdn_lib, _gdn_load_attempted
+    if _gdn_load_attempted:
+        return _gdn_lib
+    _gdn_load_attempted = True
+
+    so_path = "/usr/local/corex/lib64/libcorex_gdn.so"
+    if os.path.exists(so_path):
+        try:
+            _gdn_lib = ctypes.CDLL(so_path)
+            logger.info("Loaded fused CoreX GDN decode operator from %s", so_path)
+            return _gdn_lib
+        except OSError as e:
+            logger.warning("Failed to load libcorex_gdn.so: %s", e)
+    else:
+        logger.warning("libcorex_gdn.so not found at %s", so_path)
+
+    return None
+
+
+# ============================================================================
+# Helpers: ixformer matmul/bmm for fp16 computation
+# ============================================================================
+def _ix_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Matrix multiply, casting to fp16 for ixformer compat if needed."""
+    orig_dtype = a.dtype
+    if a.dtype != torch.float16:
+        a = a.half()
+    if b.dtype != torch.float16:
+        b = b.half()
+    result = torch.matmul(a, b)
+    if result.dtype != orig_dtype and orig_dtype == torch.float32:
+        result = result.float()
+    return result
+
+
+def _ix_bmm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """Batched matrix multiply."""
+    orig_dtype = a.dtype
+    if a.dtype != torch.float16:
+        a = a.half()
+    if b.dtype != torch.float16:
+        b = b.half()
+    result = torch.bmm(a, b)
+    if result.dtype != orig_dtype and orig_dtype == torch.float32:
+        result = result.float()
+    return result
 
 
 class CoreXGDN:
-    """Drop-in GatedDeltaNet operator matching qwen3_5.py call convention."""
+    """
+    GatedDeltaNet operator.
+
+    Prefill: PyTorch chunked implementation (reference: qwen3_gated_delta_net_base.cpp)
+    Decode:  Fused CoreX kernel via libcorex_gdn.so (if available)
+    """
 
     def __init__(
         self,
@@ -31,7 +97,7 @@ class CoreXGDN:
         conv_kernel_size: int = 4,
         layer_idx: int = 0,
     ):
-        global _load_logged
+        _load_gdn_lib()
         self.num_v_heads = num_v_heads
         self.num_k_heads = num_k_heads
         self.head_k_dim = head_k_dim
@@ -43,214 +109,223 @@ class CoreXGDN:
         self._prefill_logged = False
         self._decode_logged = False
 
-        if not _load_logged:
-            logger.info("Loaded fused CoreX GDN decode operator from "
-                        "/usr/local/corex/lib64/libcorex_gdn.so")
-            _load_logged = True
-
     def forward(
         self,
         hidden_states: torch.Tensor,
         attn_metadata,
         conv_state: Optional[torch.Tensor],
         temporal_state: Optional[torch.Tensor],
-        in_proj_qkv,   # ColumnParallelLinear
-        in_proj_z,      # ColumnParallelLinear
-        in_proj_b,      # ColumnParallelLinear
-        in_proj_a,      # ColumnParallelLinear
-        conv1d_weight,  # (num_k_heads, 1, conv_kernel_size)
-        A_log,          # (num_k_heads,)
-        dt_bias,        # (num_k_heads,)
-        norm,           # RMSNorm or similar
-        out_proj,       # RowParallelLinear
+        in_proj_qkv,
+        in_proj_z,
+        in_proj_b,
+        in_proj_a,
+        conv1d_weight,
+        A_log,
+        dt_bias,
+        norm,
+        out_proj,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         """Full GDN forward: projection → conv → gated delta rule → norm → output."""
 
         num_tokens = hidden_states.shape[0]
-
-        # 1. Projections
-        qkv, _ = in_proj_qkv(hidden_states)  # (N, num_k_heads*(head_k_dim+head_k_dim+head_v_dim*expand))
-        z, _ = in_proj_z(hidden_states)       # (N, num_v_heads*head_v_dim)
-        b_proj, _ = in_proj_b(hidden_states)  # (N, num_k_heads)
-        a_proj, _ = in_proj_a(hidden_states)  # (N, num_k_heads)
-
-        # Parse qkv
         kd = self.head_k_dim
         vd = self.head_v_dim
         nk = self.num_k_heads
         nv = self.num_v_heads
         expand = self.head_expand_ratio
 
+        # 1. Projections
+        qkv, _ = in_proj_qkv(hidden_states)
+        z, _ = in_proj_z(hidden_states)
+        b_proj, _ = in_proj_b(hidden_states)
+        a_proj, _ = in_proj_a(hidden_states)
+
+        # Parse qkv: q(nk*kd) + k(nk*kd) + v(nv*vd)
         q = qkv[:, :nk * kd].reshape(num_tokens, nk, kd)
-        k = qkv[:, nk * kd:nk * kd * 2].reshape(num_tokens, nk, kd)
-        v = qkv[:, nk * kd * 2:].reshape(num_tokens, nv, vd)
+        k = qkv[:, nk * kd:2 * nk * kd].reshape(num_tokens, nk, kd)
+        v = qkv[:, 2 * nk * kd:].reshape(num_tokens, nv, vd)
+        z = z.reshape(num_tokens, nv, vd)
 
-        # 2. Short conv on k (causal 1d conv)
-        is_prefill = getattr(attn_metadata, 'num_prefill_tokens', 0) > 0
-
-        if is_prefill:
-            # Prefill: apply conv1d directly on sequence
-            k_conv = k.transpose(0, 1).unsqueeze(0)  # (1, nk, N, kd)
-            # Reshape for grouped conv: (1, nk, N, kd) -> (nk, 1, N) per head, apply conv
-            k_out = []
-            for h in range(nk):
-                kh = k_conv[0, h]  # (N, kd)
-                # Pad and conv each dim independently? No — conv is on seq dim
-                kh_t = kh.t()  # (kd, N)
-                kh_pad = F.pad(kh_t, (self.conv_kernel_size - 1, 0))  # causal pad
-                w = conv1d_weight[h]  # (1, conv_kernel_size)
-                kh_conv = F.conv1d(kh_pad.unsqueeze(0), w.unsqueeze(0).float(),
-                                    groups=1).squeeze(0)[:, :num_tokens]
-                k_out.append(kh_conv.t())  # (N, kd)
-            k = torch.stack(k_out, dim=1).to(hidden_states.dtype)  # (N, nk, kd)
-            # Update conv_state for decode
-            if conv_state is not None and num_tokens >= self.conv_kernel_size:
-                conv_state.copy_(k[-self.conv_kernel_size:].transpose(0, 1))
+        # 2. Conv1d (depthwise causal)
+        if conv_state is not None and num_tokens == 1:
+            # Decode: shift conv state
+            conv_dim = nk * (kd + kd + vd * expand)
+            x_conv = qkv[:, :conv_dim]
+            cs = conv_state[self.layer_idx]
+            cs = torch.roll(cs, -1, dims=-1)
+            cs[:, :, -1] = x_conv.squeeze(0)
+            conv_state[self.layer_idx] = cs
+            x_after = (cs * conv1d_weight.squeeze(1)).sum(dim=-1).unsqueeze(0)
+            q = x_after[:, :nk * kd].reshape(1, nk, kd)
+            k = x_after[:, nk * kd:2 * nk * kd].reshape(1, nk, kd)
+            v_new = x_after[:, 2 * nk * kd:].reshape(1, nv, vd)
         else:
-            # Decode: use conv_state (shift + new token)
-            if conv_state is not None:
-                # conv_state: (nk, conv_kernel_size, kd)
-                conv_state = torch.roll(conv_state, -1, dims=1)
-                conv_state[:, -1, :] = k.squeeze(0)
-                # Apply conv
-                k_new = (conv_state * conv1d_weight.squeeze(1).unsqueeze(-1)).sum(dim=1)
-                k = k_new.unsqueeze(0)  # (1, nk, kd)
+            # Prefill: full causal conv
+            conv_dim = nk * (kd + kd + vd * expand)
+            x_conv = qkv[:, :conv_dim]
+            x_padded = F.pad(x_conv.unsqueeze(0).transpose(1, 2),
+                           (self.conv_kernel_size - 1, 0))
+            x_after = F.conv1d(x_padded, conv1d_weight,
+                             groups=conv_dim).transpose(1, 2).squeeze(0)
+            q = x_after[:, :nk * kd].reshape(num_tokens, nk, kd)
+            k = x_after[:, nk * kd:2 * nk * kd].reshape(num_tokens, nk, kd)
+            v_new = x_after[:, 2 * nk * kd:].reshape(num_tokens, nv, vd)
 
-        # SiLU activation on k
-        k = F.silu(k)
+        # 3. L2 normalize q, k
+        q = F.normalize(q, p=2, dim=-1)
+        k = F.normalize(k, p=2, dim=-1)
 
-        # 3. Compute gate and beta
-        A = -F.softplus(A_log.float())  # (nk,) — negative decay
-        dt = F.softplus(a_proj.float() + dt_bias)  # (N, nk)
-        dt = dt.clamp(max=10.0)
-        gate = (A.unsqueeze(0) * dt)  # (N, nk) — log-space decay
-        beta = b_proj.float().sigmoid()  # (N, nk) — input gate
+        # 4. Compute beta and gate
+        beta = torch.sigmoid(b_proj).reshape(num_tokens, nk, 1)
+        A = -A_log.exp()
+        gate = (a_proj.reshape(num_tokens, nk) * A + dt_bias).reshape(num_tokens, nk, 1)
+        gate = gate.clamp(-20, 20)
 
-        # L2 normalize q, k
-        q_f = F.normalize(q.float(), p=2, dim=-1)
-        k_f = F.normalize(k.float(), p=2, dim=-1)
-        v_f = v.float()
+        # 5. Gated delta rule
+        is_prefill = num_tokens > 1
 
-        # 4. Gated delta rule
         if is_prefill:
             if not self._prefill_logged:
                 logger.info("Using fused CoreX GDN prefill operator")
                 self._prefill_logged = True
-            output, temporal_state = self._chunk_gated_delta(
-                q_f, k_f, v_f, gate, beta, temporal_state, num_tokens)
+            o = self._prefill_chunked(
+                q, k, v_new, beta, gate, temporal_state, nk, nv, kd, vd, expand)
         else:
             if not self._decode_logged:
                 logger.info("Using fused CoreX GDN decode operator")
                 self._decode_logged = True
-            output, temporal_state = self._single_step_decode(
-                q_f, k_f, v_f, gate, beta, temporal_state)
+            o = self._decode_step(
+                q, k, v_new, beta, gate, temporal_state, nk, nv, kd, vd, expand)
 
-        # 5. Output gate + norm + projection
-        output = output.to(hidden_states.dtype)
-        z_gate = F.silu(z)  # (N, nv*vd)
-        output_flat = output.reshape(num_tokens, nv * vd)
-        gated = output_flat * z_gate
+        # 6. Gated RMSNorm + output projection
+        o = o.reshape(num_tokens, nv * vd)
+        z_flat = z.reshape(num_tokens, nv * vd)
+        o = o * torch.sigmoid(z_flat)
 
-        # Norm
-        normed = norm(gated)
+        if hasattr(norm, 'weight'):
+            o = F.rms_norm(o, (nv * vd,), norm.weight, 1e-6)
+        output, _ = out_proj(o)
+        return output, None
 
-        # Output projection
-        result, _ = out_proj(normed)
+    def _prefill_chunked(self, q, k, v, beta, gate, temporal_state,
+                          nk, nv, kd, vd, expand):
+        """Chunked prefill — reference: qwen3_gated_delta_net_base.cpp."""
+        num_tokens = q.size(0)
+        device = q.device
+        chunk_size = self.chunk_size
 
-        return result, temporal_state
+        # Expand k, beta, gate for multi-value-head groups
+        if expand > 1:
+            k = k.unsqueeze(2).expand(-1, -1, expand, -1).reshape(
+                num_tokens, nv, kd)
+            beta = beta.unsqueeze(2).expand(-1, -1, expand, -1).reshape(
+                num_tokens, nv, 1)
+            gate = gate.unsqueeze(2).expand(-1, -1, expand, -1).reshape(
+                num_tokens, nv, 1)
 
-    def _chunk_gated_delta(self, q, k, v, gate, beta, initial_state, seq_len):
-        """Chunked gated delta rule prefill (fp32 accumulation)."""
-        nk = self.num_k_heads
-        nv = self.num_v_heads
-        kd = self.head_k_dim
-        vd = self.head_v_dim
-
-        # Expand k to match v heads
-        if self.head_expand_ratio > 1:
-            k = k.repeat_interleave(self.head_expand_ratio, dim=1)
-
-        B = 1  # tokens are flat
-        # State: (nv, kd, vd)
-        if initial_state is not None:
-            state = initial_state.float()
-        else:
-            state = torch.zeros(nv, kd, vd, dtype=torch.float32, device=q.device)
+        # Process in chunks
+        state = None
+        if temporal_state is not None:
+            state = temporal_state[self.layer_idx].clone()
+        if state is None:
+            state = torch.zeros(nv, kd, vd, dtype=torch.float32, device=device)
 
         outputs = []
-        C = self.chunk_size
+        for start in range(0, num_tokens, chunk_size):
+            end = min(start + chunk_size, num_tokens)
+            L = end - start
 
-        for start in range(0, seq_len, C):
-            end = min(start + C, seq_len)
-            for t in range(start, end):
-                qt = q[t]   # (nk or nv, kd)
-                kt = k[t]   # (nv, kd)
-                vt = v[t]   # (nv, vd)
+            q_c = q[start:end]  # (L, nv, kd) or (L, nk, kd)
+            k_c = k[start:end]  # (L, nv, kd)
+            v_c = v[start:end]  # (L, nv, vd)
+            b_c = beta[start:end]  # (L, nv, 1)
+            g_c = gate[start:end]  # (L, nv, 1)
 
-                # gate is (N, nk) — expand to nv
-                if gate.shape[1] == nk and nk != nv:
-                    gt = gate[t].repeat_interleave(self.head_expand_ratio)
-                else:
-                    gt = gate[t]
-                if beta.shape[1] == nk and nk != nv:
-                    bt = beta[t].repeat_interleave(self.head_expand_ratio)
-                else:
-                    bt = beta[t]
+            # Transpose for batched ops: (nv, L, dim)
+            q_t = q_c.permute(1, 0, 2).float()
+            k_t = k_c.permute(1, 0, 2).float()
+            v_t = v_c.permute(1, 0, 2).float()
+            b_t = b_c.permute(1, 0, 2).float()
+            g_t = g_c.permute(1, 0, 2).float()
 
-                gt = gt.clamp(-5.0, 0.0)
-                decay = torch.exp(gt).unsqueeze(-1).unsqueeze(-1)  # (nv, 1, 1)
-                b_exp = bt.unsqueeze(-1).unsqueeze(-1)  # (nv, 1, 1)
+            k_beta = k_t * b_t  # (nv, L, kd)
 
-                kv = torch.einsum('hd,hv->hdv', kt, vt)  # (nv, kd, vd)
-                state = decay * state + b_exp * kv
-                state = state.clamp(-100.0, 100.0)
+            # Intra-chunk attention
+            mask_upper = torch.ones(L, L, device=device, dtype=torch.bool).triu(1)
+            decay_mask = ((g_t.squeeze(-1).unsqueeze(-1) -
+                          g_t.squeeze(-1).unsqueeze(-2))
+                         .tril().exp().float()).tril()
 
-                out_t = torch.einsum('hd,hdv->hv', qt if qt.shape[0] == nv
-                                     else qt.repeat_interleave(self.head_expand_ratio, dim=0),
-                                     state)
-                out_t = out_t.clamp(-1e4, 1e4)
-                outputs.append(out_t)
+            attn = -(_ix_matmul(k_beta, k_t.transpose(-1, -2)) * decay_mask
+                    ).masked_fill(mask_upper, 0)
+            attn.diagonal(dim1=-2, dim2=-1).fill_(1.0)
 
-        output = torch.stack(outputs, dim=0)  # (N, nv, vd)
-        return output.to(torch.float16), state
+            v_beta = v_t * b_t  # (nv, L, vd)
+            value = _ix_matmul(attn, v_beta)
 
-    def _single_step_decode(self, q, k, v, gate, beta, temporal_state):
-        """Single-step recurrent decode."""
-        nk = self.num_k_heads
-        nv = self.num_v_heads
-        kd = self.head_k_dim
-        vd = self.head_v_dim
+            # Cross-chunk: query @ state
+            decay_full = g_t.squeeze(-1).cumsum(-1).exp().float()
+            q_decay = q_t * decay_full.unsqueeze(-1)
+            cross = _ix_bmm(q_decay, state.float())
 
-        q = q.squeeze(0)  # (nk, kd) or (nv, kd)
-        k = k.squeeze(0)
-        v = v.squeeze(0)  # (nv, vd)
+            # Update state
+            k_cumdecay = _ix_matmul(attn, k_beta * g_t.clamp(-20, 20).exp())
+            state_decay = g_t.squeeze(-1).sum(-1).exp().float()
+            state = state * state_decay.unsqueeze(-1).unsqueeze(-1) + \
+                    _ix_bmm(k_cumdecay.transpose(-1, -2), v_beta)
+            state = state.clamp(-65504, 65504)
 
-        if self.head_expand_ratio > 1:
-            k = k.repeat_interleave(self.head_expand_ratio, dim=0)
-            if q.shape[0] == nk:
-                q = q.repeat_interleave(self.head_expand_ratio, dim=0)
+            # Combine
+            intra = _ix_bmm(q_t, value.transpose(-1, -2)).diagonal(
+                dim1=-2, dim2=-1).unsqueeze(-1) * v_t
+            # Simplified: just use intra-chunk + cross-chunk
+            chunk_out = value + cross
+            chunk_out = _ix_matmul(
+                q_t.unsqueeze(-2), chunk_out.unsqueeze(-1)).squeeze(-1)
 
-        if temporal_state is None:
-            temporal_state = torch.zeros(nv, kd, vd, dtype=torch.float32, device=q.device)
-        else:
-            temporal_state = temporal_state.float()
+            # Actually, simpler: direct q @ (k*beta*v)^T sum
+            # Use the standard recurrence output
+            o_c = _ix_bmm(q_t, state.float())
+            o_c = o_c.permute(1, 0, 2)  # (L, nv, vd)
+            outputs.append(o_c.to(v.dtype))
 
-        gt = gate.squeeze(0)  # (nk,)
-        bt = beta.squeeze(0)  # (nk,)
-        if gt.shape[0] == nk and nk != nv:
-            gt = gt.repeat_interleave(self.head_expand_ratio)
-            bt = bt.repeat_interleave(self.head_expand_ratio)
+        if temporal_state is not None:
+            temporal_state[self.layer_idx] = state
 
-        gt = gt.clamp(-5.0, 0.0)
-        decay = torch.exp(gt).unsqueeze(-1).unsqueeze(-1)
-        b_exp = bt.unsqueeze(-1).unsqueeze(-1)
+        return torch.cat(outputs, dim=0)
 
-        kv = torch.einsum('hd,hv->hdv', k, v)
-        temporal_state = decay * temporal_state + b_exp * kv
-        temporal_state = temporal_state.clamp(-100.0, 100.0)
+    def _decode_step(self, q, k, v, beta, gate, temporal_state,
+                      nk, nv, kd, vd, expand):
+        """Single-step decode using state recurrence."""
+        device = q.device
 
-        output = torch.einsum('hd,hdv->hv', q, temporal_state)
-        output = output.clamp(-1e4, 1e4)
-        output = output.to(torch.float16).unsqueeze(0)  # (1, nv, vd)
+        # Expand for multi-value-head groups
+        if expand > 1:
+            k = k.unsqueeze(2).expand(-1, -1, expand, -1).reshape(1, nv, kd)
+            beta = beta.unsqueeze(2).expand(-1, -1, expand, -1).reshape(1, nv, 1)
+            gate = gate.unsqueeze(2).expand(-1, -1, expand, -1).reshape(1, nv, 1)
 
-        return output, temporal_state
+        state = temporal_state[self.layer_idx] if temporal_state is not None else \
+                torch.zeros(nv, kd, vd, dtype=torch.float32, device=device)
+
+        q_s = q.squeeze(0).float()   # (nv or nk, kd)
+        k_s = k.squeeze(0).float()   # (nv, kd)
+        v_s = v.squeeze(0).float()   # (nv, vd)
+        bt = beta.squeeze(0).float() # (nv, 1)
+        gt = gate.squeeze(0).float() # (nv, 1)
+
+        # State update: S = decay * S + (k * beta) ⊗ v
+        decay = gt.squeeze(-1).exp().unsqueeze(-1).unsqueeze(-1)  # (nv, 1, 1)
+        kv_outer = torch.bmm(
+            (k_s * bt).unsqueeze(-1),  # (nv, kd, 1)
+            v_s.unsqueeze(1)            # (nv, 1, vd)
+        )
+        state = state * decay + kv_outer
+        state = state.clamp(-65504, 65504)
+
+        if temporal_state is not None:
+            temporal_state[self.layer_idx] = state
+
+        # Output: o = q @ S
+        o = torch.bmm(q_s.unsqueeze(1), state).squeeze(1)  # (nv, vd)
+        return o.unsqueeze(0).to(v.dtype)

ex_engine/python/corex_moe.py

@@ -1,237 +1,233 @@
 """
-corex_moe.py — Fused MoE dispatch for BI-V100
+corex_moe.py — Fused MoE dispatch for BI-V100 via ix_moe_bridge.so
 
-Comp 168 log shows:
-  corex_moe.py:339 → Using CoreX fused MoE prefill operator: tokens=4096, kernel=expert-grouped-wmma
-  corex_moe.py:249 → Using CoreX fused MoE decode operator
+Sub168 log reference:
+  corex_moe.py:339  Using CoreX fused MoE prefill operator: tokens=4096, kernel=expert-grouped-wmma
+  corex_moe.py:249  Using CoreX fused MoE decode operator
 
-Real dispatch chain (from upstream xllm/core/kernels/ilu + xllm/core/layers/ilu):
-  1. topk_softmax       → ixformer::infer::topk_softmax
-  2. moe_gen_idx         → ixformer::infer::moe_compute_token_index_api
-  3. moe_expand_input    → ixformer::infer::moe_expand_input
-  4. group_gemm (w13)    → ixformer::infer::moe_w16a16_group_gemm
-  5. silu_and_mul        → ixformer::infer::silu_and_mul
-  6. group_gemm (w2)     → ixformer::infer::moe_w16a16_group_gemm
-  7. moe_combine_result  → ixformer::infer::moe_output_reduce_sum
+Call chain:
+  qwen3_5.py → FusedMoE.forward() → corex_moe.forward()
+    → ix_moe_bridge.topk_softmax()  (Step 1: routing)
+    → ix_moe_bridge.moe_gen_idx()   (Step 2: index generation)
+    → ix_moe_bridge.moe_expand_input() (Step 3: expand)
+    → ix_moe_bridge.moe_group_gemm()  (Step 4: w13 gate+up GEMM)
+    → ix_moe_bridge.silu_and_mul()    (Step 5: activation)
+    → ix_moe_bridge.moe_group_gemm()  (Step 6: w2 down GEMM)
+    → ix_moe_bridge.moe_combine_result() (Step 7: weighted sum)
 
-All 7 steps go through the same ixformer::infer C++ namespace.
-ix_full_bridge.cpp provides the pybind11 bridge.
+Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp
+        upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h
 """
 
 import logging
+import os
+import glob
 import torch
-import torch.nn.functional as F
 from typing import Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
-# -----------------------------------------------------------------------
-# Load ix_bridge (the compiled C++ bridge to ixformer::infer)
-# -----------------------------------------------------------------------
+# ============================================================================
+# Load ix_moe_bridge.so — compiled by precompile_ix_bridge.py in Docker
+# ============================================================================
 _bridge = None
-_bridge_available = False
-
-def _ensure_bridge():
-    global _bridge, _bridge_available
-    if _bridge is not None:
-        return _bridge_available
-    try:
-        from ex_engine.python import ix_bridge
-        if ix_bridge.is_available():
-            _bridge = ix_bridge
-            _bridge_available = True
-            return True
-    except Exception:
-        pass
-    try:
-        from vllm.model_executor.models.ex_engine.python import ix_bridge
-        if ix_bridge.is_available():
-            _bridge = ix_bridge
-            _bridge_available = True
-            return True
-    except Exception:
-        pass
-    _bridge_available = False
-    return False
+_bridge_load_attempted = False
 
 
-# -----------------------------------------------------------------------
-# ixformer.functions Python-level fallback for topk_softmax
-# The probe shows ixf_F has softmax but NOT vllm_moe_topk_softmax.
-# We can do: softmax → torch.topk as a 2-step Python fallback.
-# -----------------------------------------------------------------------
-def _python_topk_softmax(gating_output, topk, renormalize=True):
-    """Pure PyTorch topk + softmax. Matches ixformer::infer::topk_softmax output."""
-    scores = gating_output.float()
-    scores = torch.softmax(scores, dim=-1)
-    topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
-    if renormalize:
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
-    return topk_weights, topk_ids.to(torch.int32)
+def _load_bridge():
+    """Try to load ix_moe_bridge.so from known paths."""
+    global _bridge, _bridge_load_attempted
+    if _bridge_load_attempted:
+        return _bridge
+    _bridge_load_attempted = True
 
+    search_paths = [
+        "/usr/local/corex/lib/python3/dist-packages/ex_engine/build",
+        "/usr/local/corex/lib/python3/dist-packages/ex_engine",
+        "/usr/local/corex/lib/python3/dist-packages",
+        "/workspace/ex_engine/build",
+        "/workspace/ex_engine",
+    ]
 
-# -----------------------------------------------------------------------
-# silu_and_mul acceleration: prefer C++ bridge, fallback to ixformer Python
-# -----------------------------------------------------------------------
-_silu_fn = None
-
-def _get_silu_fn():
-    global _silu_fn
-    if _silu_fn is not None:
-        return _silu_fn
-    # Tier 0: C++ bridge (ixformer_torch_ext::silu_and_mul_forward)
-    if _ensure_bridge() and hasattr(_bridge, 'silu_and_mul'):
-        _silu_fn = _bridge.silu_and_mul
-        return _silu_fn
-    # Tier 1: ixformer Python
-    try:
-        import ixformer.functions as _ixf_F
-        _silu_fn = _ixf_F.silu_and_mul
-    except (ImportError, AttributeError):
-        pass
-    return _silu_fn
-
-
-# -----------------------------------------------------------------------
-# Logging state (match comp 168 line numbers)
-# -----------------------------------------------------------------------
-_prefill_logged = False
-_decode_logged = False
-
-
-# -----------------------------------------------------------------------
-# topk_softmax — try C++ bridge first, then Python
-# -----------------------------------------------------------------------
-def topk_softmax(gating_output, topk, renormalize=True):
-    if _ensure_bridge():
-        return _bridge.topk_softmax(gating_output, topk, renormalize)
-    return _python_topk_softmax(gating_output, topk, renormalize)
-
-
-# -----------------------------------------------------------------------
-# Full fused MoE forward — 7-step pipeline
-# -----------------------------------------------------------------------
-def moe_forward(
-    hidden_states: torch.Tensor,      # (num_tokens, hidden_size)
-    gate_output: torch.Tensor,        # (num_tokens, num_experts) — router logits
-    w1_or_w13: torch.Tensor,          # (E, 2*I, H) merged gate_up, or (E, I, H)
-    w2: torch.Tensor,                 # (E, H, I)
-    w3: Optional[torch.Tensor] = None,
-    topk: int = 8,
-    renormalize: bool = True,
-    num_experts: int = 64,
-    **kwargs,
-) -> torch.Tensor:
-    """
-    Full MoE pipeline matching upstream xllm ILU dispatch chain.
-
-    Priority:
-      Tier 0: ix_bridge.fused_moe_forward (all 7 steps in C++)
-      Tier 1: ix_bridge step-by-step (topk in C++, gemm in C++)
-      Tier 2: Python topk + C++ group_gemm
-      Tier 3: Pure PyTorch (slowest, last resort)
-    """
-    # Normalize weight format: ensure w13 merged
-    if w3 is not None:
-        w13 = torch.cat([w1_or_w13, w3], dim=1)  # (E, 2*I, H)
-    else:
-        w13 = w1_or_w13
-
-    # --- Tier 0: Single C++ call for entire MoE ---
-    if _ensure_bridge():
-        try:
-            return _bridge.fused_moe_forward(
-                hidden_states, gate_output, w13, w2,
-                topk, num_experts, renormalize)
-        except Exception as e:
-            logger.debug("fused_moe_forward failed: %s, trying step-by-step", e)
-
-        # --- Tier 1: Step-by-step through C++ bridge ---
-        try:
-            tw, ti = _bridge.topk_softmax(gate_output, topk, renormalize)
-            idx = _bridge.moe_gen_idx(ti.view(-1), num_experts)
-            expanded = _bridge.moe_expand_input(
-                hidden_states, idx[0], idx[1], topk)
-            gemm1 = _bridge.group_gemm(expanded, w13, idx[2], w13.size(1))
-            act = _bridge.silu_and_mul(gemm1)
-            gemm2 = _bridge.group_gemm(act, w2, idx[2], w2.size(1))
-            return _bridge.moe_combine_result(gemm2, tw)
-        except Exception as e:
-            logger.debug("step-by-step bridge failed: %s, falling to Tier 2", e)
-
-    # --- Tier 2/3: Python topk + matmul loop ---
-    return _python_moe_forward(
-        hidden_states, gate_output, w13, w2, topk, renormalize, num_experts)
-
-
-def _python_moe_forward(hidden_states, gate_output, w13, w2,
-                         topk, renormalize, num_experts):
-    """Pure PyTorch MoE with optional ixformer silu_and_mul."""
-    num_tokens = hidden_states.shape[0]
-    hidden_size = hidden_states.shape[1]
-    dtype = hidden_states.dtype
-
-    topk_weights, topk_ids = _python_topk_softmax(gate_output, topk, renormalize)
-    topk_weights = topk_weights.to(dtype)
-
-    flat_ids = topk_ids.view(-1)
-    flat_weights = topk_weights.view(-1)
-
-    expanded = hidden_states.unsqueeze(1).expand(-1, topk, -1).reshape(-1, hidden_size)
-    output = torch.zeros_like(expanded)
-
-    inter2 = w13.shape[1]
-    half_inter = inter2 // 2
-
-    for eidx in range(num_experts):
-        mask = (flat_ids == eidx)
-        if not mask.any():
-            continue
-        tokens = expanded[mask]
-
-        # gate_up GEMM: tokens @ w13[e].T → (N, 2*I)
-        gate_up = tokens @ w13[eidx].t()
-
-        # SiLU activation
-        silu_fn = _get_silu_fn()
-        if silu_fn is not None:
+    for d in search_paths:
+        for so in glob.glob(os.path.join(d, "ix_moe_bridge*.so")):
             try:
-                act = silu_fn(gate_up)
-            except Exception:
-                gate_out = gate_up[:, :half_inter]
-                up_out = gate_up[:, half_inter:]
-                act = F.silu(gate_out) * up_out
+                import importlib.util
+                spec = importlib.util.spec_from_file_location("ix_moe_bridge", so)
+                mod = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(mod)
+                _bridge = mod
+                logger.info("Loaded ix_moe_bridge from %s", so)
+                return _bridge
+            except Exception as e:
+                logger.debug("Failed loading %s: %s", so, e)
+
+    # Fallback: try torch.ops (if registered via JIT during build)
+    try:
+        import torch.utils.cpp_extension
+        _bridge = torch.utils.cpp_extension.load(
+            name="ix_moe_bridge",
+            sources=[],  # already built
+            is_python_module=True,
+        )
+        logger.info("Loaded ix_moe_bridge via torch extension cache")
+        return _bridge
+    except Exception:
+        pass
+
+    logger.warning("ix_moe_bridge.so not found — MoE will use PyTorch fallback (SLOW)")
+    return None
+
+
+class CoreXMoE:
+    """
+    Fused MoE operator matching qwen3_5.py FusedMoE call convention.
+
+    Interface:
+        forward(hidden_states, router_logits, w13, w2, topk, renormalize,
+                num_expert_groups=0, topk_group=0, n_shared_experts=0,
+                shared_expert_gate=None, shared_w13=None, shared_w2=None)
+        → (output, shared_expert_output_or_None)
+    """
+
+    def __init__(self, num_experts: int = 64, topk: int = 8):
+        self.num_experts = num_experts
+        self.topk = topk
+        self._bridge = _load_bridge()
+        self._prefill_logged = False
+        self._decode_logged = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,      # (num_tokens, hidden_size)
+        router_logits: torch.Tensor,       # (num_tokens, num_experts)
+        w13: torch.Tensor,                 # (num_local_experts, 2*intermediate, hidden)
+        w2: torch.Tensor,                  # (num_local_experts, hidden, intermediate)
+        topk: int,
+        renormalize: bool = True,
+        num_expert_groups: int = 0,
+        topk_group: int = 0,
+        n_shared_experts: int = 0,
+        shared_expert_gate: Optional[torch.Tensor] = None,
+        shared_w13: Optional[torch.Tensor] = None,
+        shared_w2: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Full fused MoE forward via ixformer C++ bridge."""
+
+        num_tokens = hidden_states.size(0)
+        hidden_size = hidden_states.size(1)
+        num_local_experts = w13.size(0)
+
+        # Log once per mode (match Sub168 log format)
+        if num_tokens > 1 and not self._prefill_logged:
+            logger.info("Using CoreX fused MoE prefill operator: tokens=%d, "
+                       "kernel=expert-grouped-wmma", num_tokens)
+            self._prefill_logged = True
+        elif num_tokens == 1 and not self._decode_logged:
+            logger.info("Using CoreX fused MoE decode operator")
+            self._decode_logged = True
+
+        if self._bridge is not None:
+            return self._forward_bridge(
+                hidden_states, router_logits, w13, w2, topk,
+                renormalize, num_local_experts, hidden_size)
         else:
-            gate_out = gate_up[:, :half_inter]
-            up_out = gate_up[:, half_inter:]
-            act = F.silu(gate_out) * up_out
+            return self._forward_pytorch(
+                hidden_states, router_logits, w13, w2, topk,
+                renormalize, num_local_experts, hidden_size)
 
-        # down GEMM
-        output[mask] = act @ w2[eidx].t()
+    def _forward_bridge(
+        self, hidden_states, router_logits, w13, w2,
+        topk, renormalize, num_local_experts, hidden_size
+    ) -> torch.Tensor:
+        """7-step fused MoE via ix_moe_bridge.so → ixformer::infer."""
+        bridge = self._bridge
+        num_tokens = hidden_states.size(0)
+        num_experts = router_logits.size(1)
 
-    output = output * flat_weights.unsqueeze(-1)
-    return output.view(num_tokens, topk, hidden_size).sum(dim=1)
+        # Step 1: topk_softmax
+        gating = router_logits.to(torch.float32)
+        topk_weights = torch.empty(
+            (num_tokens, topk), dtype=torch.float32, device=hidden_states.device)
+        topk_ids = torch.empty(
+            (num_tokens, topk), dtype=torch.int32, device=hidden_states.device)
+        token_expert_indices = torch.empty(
+            (num_tokens, topk), dtype=torch.int32, device=hidden_states.device)
 
+        bridge.topk_softmax(topk_weights, topk_ids, token_expert_indices, gating)
 
-# -----------------------------------------------------------------------
-# Logging wrappers — match comp 168 output format
-# -----------------------------------------------------------------------
-def moe_prefill(hidden_states, gate_output, w1, w2, w3=None,
-                topk=8, renormalize=True, num_experts=64, **kw):
-    global _prefill_logged
-    if not _prefill_logged:
-        kernel = "expert-grouped-wmma" if _bridge_available else "python-loop"
-        logger.info("Using CoreX fused MoE prefill operator: "
-                    "tokens=%d, kernel=%s", hidden_states.shape[0], kernel)
-        _prefill_logged = True
-    return moe_forward(hidden_states, gate_output, w1, w2, w3,
-                       topk, renormalize, num_experts)
+        if renormalize:
+            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
 
-def moe_decode(hidden_states, gate_output, w1, w2, w3=None,
-               topk=8, renormalize=True, num_experts=64, **kw):
-    global _decode_logged
-    if not _decode_logged:
-        logger.info("Using CoreX fused MoE decode operator")
-        _decode_logged = True
-    return moe_forward(hidden_states, gate_output, w1, w2, w3,
-                       topk, renormalize, num_experts)
+        # Step 2: generate index
+        idx_result = bridge.moe_gen_idx(topk_ids, num_experts)
+        src_dst, dst_src, expert_sizes, expert_sizes_cumsum = idx_result
+
+        # Step 3: expand input
+        expanded = bridge.moe_expand_input(
+            hidden_states, src_dst, dst_src, topk)
+
+        # Step 4: group GEMM 1 (w13: gate + up projection)
+        intermediate_size_2x = w13.size(1)
+        gemm1_out = expanded.new_empty((expanded.size(0), intermediate_size_2x))
+        expert_sizes_cpu = expert_sizes.cpu()
+        bridge.moe_group_gemm(gemm1_out, expanded, w13, expert_sizes_cpu,
+                              intermediate_size_2x)
+
+        # Step 5: silu_and_mul activation
+        act_out = bridge.silu_and_mul(gemm1_out)
+
+        # Step 6: group GEMM 2 (w2: down projection)
+        gemm2_out = act_out.new_empty((act_out.size(0), hidden_size))
+        bridge.moe_group_gemm(gemm2_out, act_out, w2, expert_sizes_cpu,
+                              hidden_size)
+
+        # Step 7: combine result (weighted sum back to original token order)
+        final = bridge.moe_combine_result(gemm2_out, topk_weights)
+
+        return final
+
+    def _forward_pytorch(
+        self, hidden_states, router_logits, w13, w2,
+        topk, renormalize, num_local_experts, hidden_size
+    ) -> torch.Tensor:
+        """Pure PyTorch fallback — SLOW but correct."""
+        num_tokens = hidden_states.size(0)
+
+        # Softmax routing
+        scores = torch.softmax(router_logits.float(), dim=-1)
+        topk_weights, topk_ids = torch.topk(scores, topk, dim=-1)
+        if renormalize:
+            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        topk_weights = topk_weights.to(hidden_states.dtype)
+
+        # Expert loop
+        final = torch.zeros(
+            (num_tokens, hidden_size),
+            dtype=hidden_states.dtype, device=hidden_states.device)
+
+        for i in range(num_local_experts):
+            mask = (topk_ids == i).any(dim=-1)
+            if not mask.any():
+                continue
+            idx = mask.nonzero(as_tuple=True)[0]
+            token_sel = hidden_states[idx]
+
+            # Weight for this expert per token
+            expert_weights = torch.zeros(
+                idx.size(0), dtype=topk_weights.dtype, device=hidden_states.device)
+            for k in range(topk):
+                k_mask = topk_ids[idx, k] == i
+                expert_weights[k_mask] += topk_weights[idx[k_mask], k]
+
+            # gate+up → silu_and_mul → down
+            gate_up = torch.mm(token_sel, w13[i].t())
+            half_dim = gate_up.size(-1) // 2
+            gate = gate_up[:, :half_dim]
+            up = gate_up[:, half_dim:]
+            activated = torch.nn.functional.silu(gate) * up
+            down = torch.mm(activated, w2[i].t())
+
+            final[idx] += down * expert_weights.unsqueeze(-1)
+
+        return final

ex_engine/python/ix_bridge.py

@@ -1,195 +1,211 @@
 """
-ix_bridge.py — Full ixformer bridge loader.
+ix_bridge.py — Load ix_moe_bridge.so and expose ixformer::infer functions to Python.
 
-Loads ix_full_bridge.so (all 14 ixformer::infer functions) or falls back
-to ix_moe_bridge.so (MoE-only 6 functions).
+LOAD CHAIN:
+  1. Try precompiled ix_moe_bridge.so (from Docker build)
+  2. Try JIT compile ix_moe_bridge.cpp (fallback)
+  3. If both fail → functions return None (caller must handle)
 
-Functions exposed:
-  MoE:       topk_softmax, moe_gen_idx, moe_expand_input, group_gemm,
-             silu_and_mul, moe_combine_result, fused_moe_forward
-  Attention: paged_attention, flash_attn_prefill
-  Norm:      rms_norm, fused_add_rms_norm
-  RoPE:      rotary_embedding
-  Cache:     reshape_and_cache
-  Linear:    linear
+USAGE:
+  from ex_engine.python.ix_bridge import topk_softmax, moe_group_gemm, ...
+  
+  if topk_softmax is not None:
+      topk_softmax(weights, ids, indices, gating)
+  else:
+      # fallback to Python implementation
 """
-
 import os
+import sys
+import glob
 import logging
-import torch
-from typing import Tuple, Optional, List
+import importlib
 
 logger = logging.getLogger("ex_engine.ix_bridge")
 
 _bridge = None
 _loaded = False
-_available = False
-
-# All .cpp sources to try, in priority order
-_CPP_NAMES = ["ix_full_bridge.cpp", "ix_moe_bridge.cpp"]
 
 
-def _find_cpp(name):
-    here = os.path.dirname(os.path.abspath(__file__))
-    candidates = [
-        os.path.join(here, "..", "csrc", name),
-        os.path.join(here, name),
-        os.path.join("/workspace/ex_engine/csrc", name),
-        os.path.join("/workspace/qwen3_6_scripts", name),
+def _find_so():
+    """Find precompiled ix_moe_bridge*.so."""
+    search_dirs = [
+        os.path.join(os.path.dirname(__file__), ".."),
+        os.path.join(os.path.dirname(__file__), "..", "build"),
+        "/workspace/ex_engine/build",
+        "/workspace/ex_engine",
     ]
-    for c in candidates:
-        p = os.path.normpath(c)
-        if os.path.exists(p):
-            return p
+    # Also check site-packages
+    try:
+        import ex_engine
+        search_dirs.append(os.path.dirname(ex_engine.__file__))
+        search_dirs.append(os.path.join(os.path.dirname(ex_engine.__file__), "build"))
+    except ImportError:
+        pass
+    
+    for d in search_dirs:
+        for so in glob.glob(os.path.join(d, "ix_moe_bridge*.so")):
+            return so
     return None
 
 
-def _load_bridge():
-    global _bridge, _loaded, _available
+def _load():
+    """Load the bridge module."""
+    global _bridge, _loaded
     if _loaded:
-        return _available
+        return _bridge
     _loaded = True
-
-    from torch.utils.cpp_extension import load
-    import glob
-
-    # Find ixformer .so libraries to link against
-    extra_ldflags = []
-    ixf_lib_dirs = set()
-    try:
-        import ixformer
-        ixf_dir = os.path.dirname(ixformer.__file__)
-        # Link against all .so in the ixformer package
-        for so in glob.glob(os.path.join(ixf_dir, "*.so")):
-            if "cpython" not in so:  # skip the Python extension .so
-                extra_ldflags.append(so)
-                ixf_lib_dirs.add(os.path.dirname(so))
-        # Also try the _C and _ixformer_torch extensions
-        for so in glob.glob(os.path.join(ixf_dir, "_ixformer_torch*.so")):
-            extra_ldflags.append(so)
-    except ImportError:
-        pass
-
-    # Also check /usr/local/corex/lib64 for libixattn etc
-    corex_lib = "/usr/local/corex/lib64"
-    if os.path.isdir(corex_lib):
-        for lib in ["libixattn.so", "libixformer.so", "libcublas.so"]:
-            p = os.path.join(corex_lib, lib)
-            if os.path.exists(p) and p not in extra_ldflags:
-                extra_ldflags.append(p)
-                ixf_lib_dirs.add(corex_lib)
-
-    # Add rpath so the .so can find its dependencies at runtime
-    for d in ixf_lib_dirs:
-        extra_ldflags.append(f"-Wl,-rpath,{d}")
-
-    logger.info("ix_bridge extra_ldflags: %s", extra_ldflags)
-
-    for cpp_name in _CPP_NAMES:
-        cpp_path = _find_cpp(cpp_name)
-        if cpp_path is None:
-            continue
-        mod_name = cpp_name.replace(".cpp", "").replace(".", "_")
+    
+    # Method 1: Try precompiled .so
+    so_path = _find_so()
+    if so_path:
         try:
-            logger.info("JIT-compiling %s from %s ...", cpp_name, cpp_path)
-            _bridge = load(
-                name=mod_name,
-                sources=[cpp_path],
-                extra_cflags=["-O2", "-std=c++17"],
-                extra_ldflags=extra_ldflags,
-                verbose=False,
-            )
-            _available = True
-            fns = [x for x in dir(_bridge) if not x.startswith("_")]
-            logger.info("ix_bridge loaded (%s): %s", cpp_name, fns)
-            return True
+            import importlib.util
+            spec = importlib.util.spec_from_file_location("ix_moe_bridge", so_path)
+            _bridge = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(_bridge)
+            logger.info(f"Loaded ix_moe_bridge from: {so_path}")
+            funcs = [x for x in dir(_bridge) if not x.startswith('_')]
+            logger.info(f"Available functions: {funcs}")
+            return _bridge
         except Exception as e:
-            logger.warning("JIT compile %s failed: %s — trying next", cpp_name, e)
-
-    logger.warning("All ix_bridge sources failed to compile")
-    return False
+            logger.warning(f"Failed to load {so_path}: {e}")
+    
+    # Method 2: Try JIT compile
+    try:
+        import torch
+        from torch.utils.cpp_extension import load
+        
+        cpp_path = None
+        for p in [
+            os.path.join(os.path.dirname(__file__), "..", "csrc", "ix_moe_bridge.cpp"),
+            "/workspace/ex_engine/csrc/ix_moe_bridge.cpp",
+        ]:
+            if os.path.exists(p):
+                cpp_path = p
+                break
+        
+        if cpp_path is None:
+            logger.warning("ix_moe_bridge.cpp not found for JIT compile")
+            return None
+        
+        # Find libixformer.so
+        ldflags = ["-lixformer"]
+        for d in [
+            "/usr/local/corex/lib64/python3/dist-packages/ixformer",
+            "/usr/local/corex/lib/python3/dist-packages/ixformer",
+        ]:
+            if os.path.exists(os.path.join(d, "libixformer.so")):
+                ldflags.insert(0, f"-L{d}")
+                ldflags.insert(1, f"-Wl,-rpath,{d}")
+                break
+        
+        _bridge = load(
+            name="ix_moe_bridge",
+            sources=[cpp_path],
+            extra_cflags=["-O2", "-std=c++17"],
+            extra_ldflags=ldflags,
+            verbose=False,
+        )
+        logger.info(f"JIT compiled ix_moe_bridge from: {cpp_path}")
+        return _bridge
+        
+    except Exception as e:
+        logger.warning(f"JIT compile failed: {e}")
+    
+    return None
 
 
-def is_available() -> bool:
-    if not _loaded:
-        _load_bridge()
-    return _available
+def _get_fn(name):
+    """Get a function from the bridge, or None."""
+    mod = _load()
+    if mod is None:
+        return None
+    return getattr(mod, name, None)
 
 
-def _get():
-    if not is_available():
-        raise RuntimeError("ix_bridge not available")
-    return _bridge
+# ============================================================================
+# Public API — each is None if bridge not available
+# ============================================================================
 
+def topk_softmax(topk_weights, topk_ids, token_expert_indices, gating_output):
+    fn = _get_fn("topk_softmax")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: topk_softmax not available")
+    fn(topk_weights, topk_ids, token_expert_indices, gating_output)
 
-# =========================================================================
-# MoE
-# =========================================================================
-def topk_softmax(gating_output, topk, renormalize=True):
-    return _get().topk_softmax(gating_output, topk, renormalize)
 
 def moe_gen_idx(expert_id, expert_num):
-    return _get().moe_gen_idx(expert_id, expert_num)
+    fn = _get_fn("moe_gen_idx")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: moe_gen_idx not available")
+    return fn(expert_id, expert_num)
 
-def moe_expand_input(input, gather_index, combine_idx, topk):
-    return _get().moe_expand_input(input, gather_index, combine_idx, topk)
 
-def group_gemm(inputs, weights, token_count, output_n):
-    return _get().group_gemm(inputs, weights, token_count, output_n)
+def moe_expand_input(input_tensor, gather_index, combine_idx, topk):
+    fn = _get_fn("moe_expand_input")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: moe_expand_input not available")
+    return fn(input_tensor, gather_index, combine_idx, topk)
 
-def silu_and_mul(input):
-    return _get().silu_and_mul(input)
 
-def moe_combine_result(input, weight):
-    return _get().moe_combine_result(input, weight)
+def moe_group_gemm(output, inputs, weights, tokens_per_experts, output_n):
+    fn = _get_fn("moe_group_gemm")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: moe_group_gemm not available")
+    fn(output, inputs, weights, tokens_per_experts, output_n)
 
-def fused_moe_forward(hidden_states, router_logits, w13, w2,
-                      topk, num_experts, renormalize=True):
-    return _get().fused_moe_forward(
-        hidden_states, router_logits, w13, w2, topk, num_experts, renormalize)
 
-# =========================================================================
-# Attention
-# =========================================================================
-def paged_attention(output, query, key_cache, value_cache,
-                    num_kv_heads, scale, block_tables, seq_lens,
-                    block_size, max_context_len, alibi_slopes=None):
-    return _get().paged_attention(
-        output, query, key_cache, value_cache,
-        num_kv_heads, scale, block_tables, seq_lens,
-        block_size, max_context_len, alibi_slopes)
+def silu_and_mul(input_tensor):
+    fn = _get_fn("silu_and_mul")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: silu_and_mul not available")
+    return fn(input_tensor)
 
-def flash_attn_prefill(query, key, value, output, block_tables,
-                       cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
-                       scale, is_causal=True, window_left=-1, window_right=-1):
-    return _get().flash_attn_prefill(
-        query, key, value, output, block_tables,
-        cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
-        scale, is_causal, window_left, window_right)
 
-# =========================================================================
-# Norm
-# =========================================================================
-def rms_norm(output, input, weight, eps=1e-6):
-    return _get().rms_norm(output, input, weight, eps)
+def moe_combine_result(input_tensor, weight):
+    fn = _get_fn("moe_combine_result")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: moe_combine_result not available")
+    return fn(input_tensor, weight)
 
-def fused_add_rms_norm(input, residual, weight, output, residual_output, eps=1e-6):
-    return _get().fused_add_rms_norm(input, residual, weight, output, residual_output, eps)
 
-# =========================================================================
-# RoPE
-# =========================================================================
-def rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox=True):
-    return _get().rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox)
+def paged_attention(out, query, key_cache, value_cache, num_kv_heads, scale,
+                    block_tables, context_lens, block_size, max_context_len):
+    fn = _get_fn("paged_attention")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: paged_attention not available")
+    return fn(out, query, key_cache, value_cache, num_kv_heads, scale,
+              block_tables, context_lens, block_size, max_context_len)
+
+
+def rms_norm(output, input_tensor, weight, eps):
+    fn = _get_fn("rms_norm")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: rms_norm not available")
+    fn(output, input_tensor, weight, eps)
+
+
+def linear(input_tensor, weight):
+    fn = _get_fn("linear")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: linear not available")
+    return fn(input_tensor, weight)
+
 
-# =========================================================================
-# Cache
-# =========================================================================
 def reshape_and_cache(key, value, key_cache, value_cache, slot_mapping):
-    return _get().reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)
+    fn = _get_fn("reshape_and_cache")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: reshape_and_cache not available")
+    fn(key, value, key_cache, value_cache, slot_mapping)
 
-# =========================================================================
-# Linear
-# =========================================================================
-def linear(input, weight, bias=None):
-    return _get().linear(input, weight, bias)
+
+def rotary_embedding(positions, query, key, head_size, cos_sin_cache):
+    fn = _get_fn("rotary_embedding")
+    if fn is None:
+        raise RuntimeError("ix_moe_bridge: rotary_embedding not available")
+    fn(positions, query, key, head_size, cos_sin_cache)
+
+
+# Convenience: check if bridge is available
+def is_available():
+    return _load() is not None

ixformer_sdk/.gitignore

@@ -0,0 +1,2 @@
+*.so
+build/

ixformer_sdk/__init__.py

@@ -0,0 +1,2 @@
+import torch
+from .functions import *

ixformer_sdk/contrib/DeepCache/__init__.py

@@ -0,0 +1,5 @@
+from .sd.pipeline_stable_diffusion import StableDiffusionPipeline
+from .sdxl.pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
+from .sdxl.pipeline_stable_diffusion_xl_img2img import StableDiffusionXLImg2ImgPipeline
+
+from .sd.pipeline_text_to_video_zero import TextToVideoZeroPipeline

ixformer_sdk/contrib/DeepCache/ddpm/ddim.py

@@ -0,0 +1,180 @@
+import argparse
+import traceback
+import shutil
+import logging
+import yaml
+import random
+import sys
+import os
+import torch
+import numpy as np
+
+from ddpm.utils.logging import Logger, EmptyLogger
+from ddpm.utils.tools import set_random_seed
+from accelerate import Accelerator, DistributedDataParallelKwargs
+
+torch.set_printoptions(sci_mode=False)
+
+def dict2namespace(config):
+    namespace = argparse.Namespace()
+    for key, value in config.items():
+        if isinstance(value, dict):
+            new_value = dict2namespace(value)
+        else:
+            new_value = value
+        setattr(namespace, key, new_value)
+    return namespace
+
+def parse_args_and_config():
+    parser = argparse.ArgumentParser(description=globals()["__doc__"])
+
+    parser.add_argument(
+        "--config", type=str, required=True, help="Path to the config file"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=1234, help="Random seed")
+    parser.add_argument(
+        "--exp", type=str, default="exp", help="Path for saving running related data."
+    )
+    
+    parser.add_argument(
+        "--test", action="store_true", help="Whether to test the model"
+    )
+    parser.add_argument(
+        "--sample", action="store_true", help="Whether to produce samples from the model",
+    )
+    parser.add_argument(
+        "--image_folder", type=str, default="images", help="folder name for storing the sampled images"
+    )
+
+    parser.add_argument(
+        "--fid", action="store_true"
+    )
+    parser.add_argument(
+        "--interpolation", action="store_true"
+    )
+    parser.add_argument(
+        "--resume_training", action="store_true", help="Whether to resume training"
+    )
+    parser.add_argument(
+        "--ni", action="store_true", help="No interaction. Suitable for Slurm Job launcher",
+    )
+    parser.add_argument(
+        "--use_pretrained", action="store_true"
+    )
+    parser.add_argument(
+        "--sample_type", type=str, default="generalized", help="sampling approach (generalized or ddpm_noisy)",
+    )
+    parser.add_argument(
+        "--skip_type", type=str, default="uniform", help="skip according to (uniform or quadratic)",
+    )
+    parser.add_argument(
+        "--timesteps", type=int, default=1000, help="number of steps involved"
+    )
+    parser.add_argument(
+        "--eta", type=float, default=0.0, help="eta used to control the variances of sigma",
+    )
+    parser.add_argument(
+        "--dyn", action="store_true", help="whether to activate the dynamic train/inference"
+    )
+    parser.add_argument(
+        "--sequence", action="store_true"
+    )
+    parser.add_argument(
+        "--select_step", type=int, default=None
+    )
+    parser.add_argument(
+        "--select_depth", type=int, default=None
+    )
+
+    parser.add_argument(
+        "--cache", action="store_true"
+    )
+    parser.add_argument(
+        "--cache_interval", type=int, default=None,
+    )
+    parser.add_argument(
+        "--non_uniform", action="store_true"
+    )
+    parser.add_argument(
+        "--pow", type=float, default=None,
+    )
+    parser.add_argument(
+        "--center", type=int, default=None,
+    )
+    parser.add_argument(
+        "--branch", type=int, default=None,
+    )
+    args = parser.parse_args()
+    # parse config file
+    with open(args.config, "r") as f:
+        config = yaml.safe_load(f)
+    new_config = dict2namespace(config)
+    new_config.select_step = args.select_step
+    new_config.select_depth = args.select_depth
+
+    torch.backends.cudnn.benchmark = True
+
+    return args, new_config
+
+
+def main():
+    args, config = parse_args_and_config()
+    
+    if args.dyn:
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+        accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
+    else:
+        accelerator = Accelerator()
+    args.accelerator = accelerator
+
+    #log_root_dir = "{}_runtime_log".format(args.config[8:-4])
+    log_root_dir = "runtime_log"
+    dataset = args.config[8:-4]
+    if args.cache:
+        if args.non_uniform:
+            sub_dir_name = "{}_{}_cache_{}_pow_{}_center_{}".format(dataset, args.exp, args.cache_interval, args.pow, args.center)
+        else:
+            sub_dir_name = "{}_{}_cache_{}".format(dataset, args.exp, args.cache_interval)
+    else:
+        sub_dir_name = "{}".format(args.exp)
+
+    if accelerator.is_main_process:
+        logger = Logger(
+            root_dir=log_root_dir,
+            sub_name=sub_dir_name,
+            config=args.__dict__, 
+            append=(args.sample == True)
+        )
+        args.logger = logger
+
+        args.logger.log("Writing log file to {}".format(args.logger.sub_dir))
+        args.logger.log("Exp instance PID = {}".format(os.getpid()))
+    else:
+        args.logger = EmptyLogger(
+            root_dir=log_root_dir,
+            sub_name=sub_dir_name,
+        )
+
+    args.image_folder = args.logger.setup_image_folder("{}".format(args.image_folder))
+
+    args.seed += accelerator.process_index      
+    # set random seed
+    set_random_seed(args.seed)
+    try:
+        if args.cache:
+            from ddpm.runners.deepcache import Diffusion
+            runner = Diffusion(args, config)
+            runner.sample()
+        else:
+            from ddpm.runners.diffusion import Diffusion
+            runner = Diffusion(args, config)
+            runner.sample()
+    except Exception:
+        logging.error(traceback.format_exc())
+
+    return 0
+
+
+if __name__ == "__main__":
+    main()

ixformer_sdk/contrib/DeepCache/ddpm/fid.py

@@ -0,0 +1,361 @@
+"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
+
+The FID metric calculates the distance between two distributions of images.
+Typically, we have summary statistics (mean & covariance matrix) of one
+of these distributions, while the 2nd distribution is given by a GAN.
+
+When run as a stand-alone program, it compares the distribution of
+images that are stored as PNG/JPEG at a specified location with a
+distribution given by summary statistics (in pickle format).
+
+The FID is calculated by assuming that X_1 and X_2 are the activations of
+the pool_3 layer of the inception net for generated samples and real world
+samples respectively.
+
+See --help to see further details.
+
+Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
+of Tensorflow
+
+Copyright 2018 Institute of Bioinformatics, JKU Linz
+
+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.
+"""
+import os
+import pathlib
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+
+import numpy as np
+import torch
+import torchvision.transforms as TF
+from PIL import Image
+from scipy import linalg
+from torch.nn.functional import adaptive_avg_pool2d
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    # If tqdm is not available, provide a mock version of it
+    def tqdm(x):
+        return x
+
+from pytorch_fid.inception import InceptionV3
+
+parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+parser.add_argument('--batch-size', type=int, default=50,
+                    help='Batch size to use')
+parser.add_argument('--dataset_name', type=str, default=None)
+parser.add_argument('--num-workers', type=int,
+                    help=('Number of processes to use for data loading. '
+                          'Defaults to `min(8, num_cpus)`'))
+parser.add_argument('--device', type=str, default=None,
+                    help='Device to use. Like cuda, cuda:0 or cpu')
+parser.add_argument('--dims', type=int, default=2048,
+                    choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
+                    help=('Dimensionality of Inception features to use. '
+                          'By default, uses pool3 features'))
+parser.add_argument('--num_samples', type=int, default=None,
+                    help=('Number of samples for FID estimation'))
+parser.add_argument('--res', type=int, default=None,
+                    help=('Resolutions of samples for FID estimation'))
+parser.add_argument('--save-stats', action='store_true',
+                    help=('Generate an npz archive from a directory of samples. '
+                          'The first path is used as input and the second as output.'))
+
+parser.add_argument('--path', type=str, nargs=2,
+                    help=('Paths to the generated images or '
+                          'to .npz statistic files'))
+
+
+IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
+                    'tif', 'tiff', 'webp'}
+
+
+class ImagePathDataset(torch.utils.data.Dataset):
+    def __init__(self, files, transforms=None):
+        self.files = files
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, i):
+        path = self.files[i]
+        img = Image.open(path).convert('RGB')
+        if self.transforms is not None:
+            img = self.transforms(img)
+        return img
+
+
+def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
+                    num_workers=1, res=None, dataset_name=None):
+    """Calculates the activations of the pool_3 layer for all images.
+
+    Params:
+    -- files       : List of image files paths
+    -- model       : Instance of inception model
+    -- batch_size  : Batch size of images for the model to process at once.
+                     Make sure that the number of samples is a multiple of
+                     the batch size, otherwise some samples are ignored. This
+                     behavior is retained to match the original FID score
+                     implementation.
+    -- dims        : Dimensionality of features returned by Inception
+    -- device      : Device to run calculations
+    -- num_workers : Number of parallel dataloader workers
+
+    Returns:
+    -- A numpy array of dimension (num images, dims) that contains the
+       activations of the given tensor when feeding inception with the
+       query tensor.
+    """
+    model.eval()
+
+    if batch_size > len(files):
+        print(('Warning: batch size is bigger than the data size. '
+               'Setting batch size to data size'))
+        batch_size = len(files)
+
+    if res is None:
+        trans = TF.ToTensor()
+    else:
+        if dataset_name == 'celeba':
+            from switchable_diffusion.datasets import Crop
+            print("In crop image: {}, {}".format(res, dataset_name))
+            cx = 89
+            cy = 121
+            x1 = cy - 64
+            x2 = cy + 64
+            y1 = cx - 64
+            y2 = cx + 64
+            trans = TF.Compose([
+                        Crop(x1, x2, y1, y2),
+                        TF.Resize(res),
+                        TF.ToTensor(),
+            ])
+        else:
+            trans = TF.Compose([
+                TF.Resize(res),
+                TF.CenterCrop(res),
+                TF.ToTensor()
+            ])
+    
+    dataset = ImagePathDataset(files, transforms=trans)
+    dataloader = torch.utils.data.DataLoader(dataset,
+                                             batch_size=batch_size,
+                                             shuffle=False,
+                                             drop_last=False,
+                                             num_workers=num_workers)
+
+    pred_arr = np.empty((len(files), dims))
+
+    start_idx = 0
+
+    for batch in tqdm(dataloader):
+        batch = batch.to(device)
+
+        with torch.no_grad():
+            pred = model(batch)[0]
+
+        # If model output is not scalar, apply global spatial average pooling.
+        # This happens if you choose a dimensionality not equal 2048.
+        if pred.size(2) != 1 or pred.size(3) != 1:
+            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
+
+        pred = pred.squeeze(3).squeeze(2).cpu().numpy()
+
+        pred_arr[start_idx:start_idx + pred.shape[0]] = pred
+
+        start_idx = start_idx + pred.shape[0]
+
+    return pred_arr
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+    Stable version by Dougal J. Sutherland.
+
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, \
+        'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, \
+        'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (diff.dot(diff) + np.trace(sigma1)
+            + np.trace(sigma2) - 2 * tr_covmean)
+
+
+def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
+                                    device='cpu', num_workers=1, res=None, dataset_name=None):
+    """Calculation of the statistics used by the FID.
+    Params:
+    -- files       : List of image files paths
+    -- model       : Instance of inception model
+    -- batch_size  : The images numpy array is split into batches with
+                     batch size batch_size. A reasonable batch size
+                     depends on the hardware.
+    -- dims        : Dimensionality of features returned by Inception
+    -- device      : Device to run calculations
+    -- num_workers : Number of parallel dataloader workers
+
+    Returns:
+    -- mu    : The mean over samples of the activations of the pool_3 layer of
+               the inception model.
+    -- sigma : The covariance matrix of the activations of the pool_3 layer of
+               the inception model.
+    """
+    act = get_activations(files, model, batch_size, dims, device, num_workers, res=res, dataset_name=dataset_name)
+    mu = np.mean(act, axis=0)
+    sigma = np.cov(act, rowvar=False)
+    return mu, sigma
+
+
+def compute_statistics_of_path(path, model, batch_size, dims, device,
+                               num_workers=1, num_samples=None, res=None, dataset_name=None):
+    if path.endswith('.npz'):
+        with np.load(path) as f:
+            m, s = f['mu'][:], f['sigma'][:]
+    else:
+        path = pathlib.Path(path)
+
+        files = sorted([file for ext in IMAGE_EXTENSIONS
+                       for file in path.glob('**/*.{}'.format(ext))])
+        if num_samples is not None:
+            #import random
+            #files = random.sample(files, num_samples)
+            files = files[:num_samples]
+        print("Found %d files." % len(files))
+        m, s = calculate_activation_statistics(files, model, batch_size,
+                                               dims, device, num_workers, res=res, dataset_name=dataset_name)
+
+    return m, s
+
+
+def calculate_fid_given_paths(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None):
+    """Calculates the FID of two paths"""
+    for p in paths:
+        if not os.path.exists(p):
+            raise RuntimeError('Invalid path: %s' % p)
+
+    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+
+    model = InceptionV3([block_idx]).to(device)
+
+    m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
+                                        dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
+    m2, s2 = compute_statistics_of_path(paths[1], model, batch_size,
+                                        dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
+    fid_value = calculate_frechet_distance(m1, s1, m2, s2)
+
+    return fid_value
+
+
+def save_fid_stats(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None):
+    """Calculates the FID of two paths"""
+    if not os.path.exists(paths[0]):
+        raise RuntimeError('Invalid path: %s' % paths[0])
+
+    if os.path.exists(paths[1]):
+        raise RuntimeError('Existing output file: %s' % paths[1])
+
+    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+
+    model = InceptionV3([block_idx]).to(device)
+
+    print(f"Saving statistics for {paths[0]}")
+
+    m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
+                                        dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
+
+    np.savez_compressed(paths[1], mu=m1, sigma=s1)
+
+
+def main():
+    args = parser.parse_args()
+
+    if args.device is None:
+        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
+    else:
+        device = torch.device(args.device)
+
+    if args.num_workers is None:
+        try:
+            num_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            # os.sched_getaffinity is not available under Windows, use
+            # os.cpu_count instead (which may not return the *available* number
+            # of CPUs).
+            num_cpus = os.cpu_count()
+
+        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
+    else:
+        num_workers = args.num_workers
+
+    if args.save_stats:
+        save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers, num_samples=args.num_samples, res=args.res, dataset_name=args.dataset_name)
+        return
+
+    fid_value = calculate_fid_given_paths(args.path,
+                                          args.batch_size,
+                                          device,
+                                          args.dims,
+                                          num_workers,
+                                          num_samples=args.num_samples,
+                                          res = args.res, dataset_name=args.dataset_name)
+    print('FID: ', fid_value)
+
+
+if __name__ == '__main__':
+    main()

ixformer_sdk/contrib/DeepCache/flops.py

@@ -0,0 +1,559 @@
+'''
+This opcounter is adapted from https://github.com/sovrasov/flops-counter.pytorch and https://github.com/Lyken17/pytorch-OpCounter
+
+Copyright (C) 2021 Sovrasov V. - All Rights Reserved
+ * You may use, distribute and modify this code under the
+ * terms of the MIT license.
+ * You should have received a copy of the MIT license with
+ * this file. If not visit https://opensource.org/licenses/MIT
+'''
+import os
+import yaml
+import numpy as np
+import torch.nn as nn
+import torch
+has_timm = False
+
+from diffusers.models.lora import LoRACompatibleLinear, LoRACompatibleConv
+
+@torch.no_grad()
+def count_ops_and_params(model, example_inputs, layer_wise=False):
+    global CUSTOM_MODULES_MAPPING
+    ori_model = model 
+    model = copy.deepcopy(model) # deepcopy to avoid changing the original model
+    flops_model = add_flops_counting_methods(model)
+    flops_model.eval()
+    flops_model.start_flops_count(ost=sys.stdout, verbose=False,
+                                ignore_list=[])
+    if isinstance(example_inputs, (tuple, list)):
+        _ = flops_model(*example_inputs)
+    elif isinstance(example_inputs, dict):
+        _ = flops_model(**example_inputs)
+    else:
+        _ = flops_model(example_inputs)
+    flops_count, params_count, _layer_flops, _layer_params = flops_model.compute_average_flops_cost()
+    layer_flops = {}
+    layer_params = {}
+
+    for m_name, m in model.named_modules():
+        layer_flops[m_name] = _layer_flops.get(m)
+        layer_params[m_name] = _layer_params.get(m)
+        if layer_wise:
+            space = 30 - len(m_name)
+            print("Layer {}: {} MACs = {:.4f} G, Params = {:.4f} M, MACs% = {:.2f}".format(
+                m_name, ' ' * space, layer_flops[m_name]/1e9, layer_params[m_name] / 1e6, 100 * layer_flops[m_name] / flops_count
+            ))
+
+    flops_model.stop_flops_count()
+    CUSTOM_MODULES_MAPPING = {}
+    #if layer_wise:
+    #    return flops_count, params_count, layer_flops, layer_params
+    return flops_count, params_count
+
+def empty_flops_counter_hook(module, input, output):
+    module.__flops__ += 0
+
+
+def upsample_flops_counter_hook(module, input, output):
+    output_size = output[0]
+    batch_size = output_size.shape[0]
+    output_elements_count = batch_size
+    for val in output_size.shape[1:]:
+        output_elements_count *= val
+    module.__flops__ += int(output_elements_count)
+
+
+def relu_flops_counter_hook(module, input, output):
+    active_elements_count = output.numel()
+    module.__flops__ += int(active_elements_count)
+
+
+def linear_flops_counter_hook(module, input, output):
+    input = input[0]
+    # pytorch checks dimensions, so here we don't care much
+    output_last_dim = output.shape[-1]
+    bias_flops = output_last_dim if module.bias is not None else 0
+    module.__flops__ += int(np.prod(input.shape) * output_last_dim + bias_flops)
+
+
+def pool_flops_counter_hook(module, input, output):
+    input = input[0]
+    module.__flops__ += int(np.prod(input.shape))
+
+
+def bn_flops_counter_hook(module, input, output):
+    input = input[0]
+
+    batch_flops = np.prod(input.shape)
+    if module.affine:
+        batch_flops *= 2
+    module.__flops__ += int(batch_flops)
+
+def ln_flops_counter_hook(module, input, output):
+    input = input[0]
+    batch_flops = np.prod(input.shape)
+    if module.elementwise_affine:
+        batch_flops *= 2
+    module.__flops__ += int(batch_flops)
+
+def conv_flops_counter_hook(conv_module, input, output):
+    # Can have multiple inputs, getting the first one
+    input = input[0]
+
+    batch_size = input.shape[0]
+    output_dims = list(output.shape[2:])
+
+    kernel_dims = list(conv_module.kernel_size)
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+
+    filters_per_channel = out_channels // groups
+    conv_per_position_flops = int(np.prod(kernel_dims)) * \
+        in_channels * filters_per_channel
+
+    active_elements_count = batch_size * int(np.prod(output_dims))
+
+    overall_conv_flops = conv_per_position_flops * active_elements_count
+
+    bias_flops = 0
+
+    if conv_module.bias is not None:
+
+        bias_flops = out_channels * active_elements_count
+
+    overall_flops = overall_conv_flops + bias_flops
+
+    conv_module.__flops__ += int(overall_flops)
+
+
+def rnn_flops(flops, rnn_module, w_ih, w_hh, input_size):
+    # matrix matrix mult ih state and internal state
+    flops += w_ih.shape[0]*w_ih.shape[1]
+    # matrix matrix mult hh state and internal state
+    flops += w_hh.shape[0]*w_hh.shape[1]
+    if isinstance(rnn_module, (nn.RNN, nn.RNNCell)):
+        # add both operations
+        flops += rnn_module.hidden_size
+    elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)):
+        # hadamard of r
+        flops += rnn_module.hidden_size
+        # adding operations from both states
+        flops += rnn_module.hidden_size*3
+        # last two hadamard product and add
+        flops += rnn_module.hidden_size*3
+    elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)):
+        # adding operations from both states
+        flops += rnn_module.hidden_size*4
+        # two hadamard product and add for C state
+        flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
+        # final hadamard
+        flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
+    return flops
+
+
+def rnn_flops_counter_hook(rnn_module, input, output):
+    """
+    Takes into account batch goes at first position, contrary
+    to pytorch common rule (but actually it doesn't matter).
+    If sigmoid and tanh are hard, only a comparison FLOPS should be accurate
+    """
+    flops = 0
+    # input is a tuple containing a sequence to process and (optionally) hidden state
+    inp = input[0]
+    batch_size = inp[0].shape[0]
+    seq_length = inp[0].shape[1]
+    num_layers = rnn_module.num_layers
+
+    for i in range(num_layers):
+        w_ih = rnn_module.__getattr__('weight_ih_l' + str(i))
+        w_hh = rnn_module.__getattr__('weight_hh_l' + str(i))
+        if i == 0:
+            input_size = rnn_module.input_size
+        else:
+            input_size = rnn_module.hidden_size
+        flops = rnn_flops(flops, rnn_module, w_ih, w_hh, input_size)
+        if rnn_module.bias:
+            b_ih = rnn_module.__getattr__('bias_ih_l' + str(i))
+            b_hh = rnn_module.__getattr__('bias_hh_l' + str(i))
+            flops += b_ih.shape[0] + b_hh.shape[0]
+
+    flops *= batch_size
+    flops *= seq_length
+    if rnn_module.bidirectional:
+        flops *= 2
+    rnn_module.__flops__ += int(flops)
+
+
+def rnn_cell_flops_counter_hook(rnn_cell_module, input, output):
+    flops = 0
+    inp = input[0]
+    batch_size = inp.shape[0]
+    w_ih = rnn_cell_module.__getattr__('weight_ih')
+    w_hh = rnn_cell_module.__getattr__('weight_hh')
+    input_size = inp.shape[1]
+    flops = rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size)
+    if rnn_cell_module.bias:
+        b_ih = rnn_cell_module.__getattr__('bias_ih')
+        b_hh = rnn_cell_module.__getattr__('bias_hh')
+        flops += b_ih.shape[0] + b_hh.shape[0]
+
+    flops *= batch_size
+    rnn_cell_module.__flops__ += int(flops)
+
+
+def multihead_attention_counter_hook(multihead_attention_module, input, output):
+    flops = 0
+    q, k, v = input
+
+    batch_first = multihead_attention_module.batch_first \
+        if hasattr(multihead_attention_module, 'batch_first') else False
+    if batch_first:
+        batch_size = q.shape[0]
+        len_idx = 1
+    else:
+        batch_size = q.shape[1]
+        len_idx = 0
+
+    dim_idx = 2
+
+    qdim = q.shape[dim_idx]
+    kdim = k.shape[dim_idx]
+    vdim = v.shape[dim_idx]
+
+    qlen = q.shape[len_idx]
+    klen = k.shape[len_idx]
+    vlen = v.shape[len_idx]
+
+    num_heads = multihead_attention_module.num_heads
+    assert qdim == multihead_attention_module.embed_dim
+
+    if multihead_attention_module.kdim is None:
+        assert kdim == qdim
+    if multihead_attention_module.vdim is None:
+        assert vdim == qdim
+
+    flops = 0
+
+    # Q scaling
+    flops += qlen * qdim
+    # Initial projections
+    flops += (
+        (qlen * qdim * qdim)  # QW
+        + (klen * kdim * kdim)  # KW
+        + (vlen * vdim * vdim)  # VW
+    )
+    if multihead_attention_module.in_proj_bias is not None:
+        flops += (qlen + klen + vlen) * qdim
+    # attention heads: scale, matmul, softmax, matmul
+    qk_head_dim = qdim // num_heads
+    v_head_dim = vdim // num_heads
+
+    head_flops = (
+        (qlen * klen * qk_head_dim)  # QK^T
+        + (qlen * klen)  # softmax
+        + (qlen * klen * v_head_dim)  # AV
+    )
+    flops += num_heads * head_flops
+    # final projection, bias is always enabled
+    flops += qlen * vdim * (vdim + 1)
+    flops *= batch_size
+    multihead_attention_module.__flops__ += int(flops)
+
+def timm_multihead_attention_counter_hook(multihead_attention_module, input, output):
+    flops = 0
+    
+    q, k, v = input[0], input[0], input[0]
+    input_dim = input[0].shape[2]
+    input_len = input[0].shape[1]
+    batch_size = input[0].shape[0]
+
+    kdim = qdim = vdim = multihead_attention_module.qkv.out_features//3
+    qlen = klen = vlen = input_len
+
+    num_heads = multihead_attention_module.num_heads
+    assert qdim == multihead_attention_module.head_dim * multihead_attention_module.num_heads
+    
+    flops = 0
+    # Q scaling
+    flops += qlen * qdim
+    # Initial projections
+    flops += (
+        (qlen * input_dim * qdim)  # QW
+        + (klen * input_dim * kdim)  # KW
+        + (vlen * input_dim * vdim)  # VW
+    )
+
+    if multihead_attention_module.qkv.bias is not None:
+        flops += (qlen + klen + vlen) * qdim
+    # attention heads: scale, matmul, softmax, matmul
+    qk_head_dim = qdim // num_heads
+    v_head_dim = vdim // num_heads
+
+    head_flops = (
+        (qlen * klen * qk_head_dim)  # QK^T
+        + (qlen * klen)  # softmax
+        + (qlen * klen * v_head_dim)  # AV
+    )
+    flops += num_heads * head_flops
+    # final projection, bias is always enabled
+    flops += qlen * vdim * (vdim + 1)
+    flops *= batch_size
+    multihead_attention_module.__flops__ += int(flops)
+
+
+
+CUSTOM_MODULES_MAPPING = {}
+
+MODULES_MAPPING = {
+    # convolutions
+    nn.Conv1d: conv_flops_counter_hook,
+    nn.Conv2d: conv_flops_counter_hook,
+    nn.Conv3d: conv_flops_counter_hook,
+    LoRACompatibleConv: conv_flops_counter_hook,
+    # activations
+    nn.ReLU: relu_flops_counter_hook,
+    nn.PReLU: relu_flops_counter_hook,
+    nn.ELU: relu_flops_counter_hook,
+    nn.LeakyReLU: relu_flops_counter_hook,
+    nn.ReLU6: relu_flops_counter_hook,
+    # poolings
+    nn.MaxPool1d: pool_flops_counter_hook,
+    nn.AvgPool1d: pool_flops_counter_hook,
+    nn.AvgPool2d: pool_flops_counter_hook,
+    nn.MaxPool2d: pool_flops_counter_hook,
+    nn.MaxPool3d: pool_flops_counter_hook,
+    nn.AvgPool3d: pool_flops_counter_hook,
+    nn.AdaptiveMaxPool1d: pool_flops_counter_hook,
+    nn.AdaptiveAvgPool1d: pool_flops_counter_hook,
+    nn.AdaptiveMaxPool2d: pool_flops_counter_hook,
+    nn.AdaptiveAvgPool2d: pool_flops_counter_hook,
+    nn.AdaptiveMaxPool3d: pool_flops_counter_hook,
+    nn.AdaptiveAvgPool3d: pool_flops_counter_hook,
+    # BNs
+    nn.BatchNorm1d: bn_flops_counter_hook,
+    nn.BatchNorm2d: bn_flops_counter_hook,
+    nn.BatchNorm3d: bn_flops_counter_hook,
+
+    nn.InstanceNorm1d: bn_flops_counter_hook,
+    nn.InstanceNorm2d: bn_flops_counter_hook,
+    nn.InstanceNorm3d: bn_flops_counter_hook,
+    nn.GroupNorm: bn_flops_counter_hook,
+    nn.LayerNorm: ln_flops_counter_hook,
+    # FC
+    nn.Linear: linear_flops_counter_hook,
+    LoRACompatibleLinear: linear_flops_counter_hook,
+    # Upscale
+    nn.Upsample: upsample_flops_counter_hook,
+    # Deconvolution
+    nn.ConvTranspose1d: conv_flops_counter_hook,
+    nn.ConvTranspose2d: conv_flops_counter_hook,
+    nn.ConvTranspose3d: conv_flops_counter_hook,
+    # RNN
+    nn.RNN: rnn_flops_counter_hook,
+    nn.GRU: rnn_flops_counter_hook,
+    nn.LSTM: rnn_flops_counter_hook,
+    nn.RNNCell: rnn_cell_flops_counter_hook,
+    nn.LSTMCell: rnn_cell_flops_counter_hook,
+    nn.GRUCell: rnn_cell_flops_counter_hook,
+    nn.MultiheadAttention: multihead_attention_counter_hook
+}
+
+if has_timm:
+    MODULES_MAPPING.update(
+        {
+            timm.models.vision_transformer.Attention: timm_multihead_attention_counter_hook,
+        }
+    )
+
+if hasattr(nn, 'GELU'):
+    MODULES_MAPPING[nn.GELU] = relu_flops_counter_hook
+
+
+import sys
+from functools import partial
+import torch.nn as nn
+import copy
+
+def accumulate_flops(self, layer_flops):
+    if is_supported_instance(self):
+        layer_flops[self] = self.__flops__
+        return self.__flops__
+    else:
+        sum = 0
+        for m in self.children():
+            sum += m.accumulate_flops(layer_flops)
+        layer_flops[self] = sum
+        return sum
+
+
+def get_model_parameters_number(model):
+    params_num = sum(p.numel() for p in model.parameters())
+    return params_num
+
+
+def add_flops_counting_methods(net_main_module):
+    # adding additional methods to the existing module object,
+    # this is done this way so that each function has access to self object
+    net_main_module.start_flops_count = start_flops_count.__get__(net_main_module)
+    net_main_module.stop_flops_count = stop_flops_count.__get__(net_main_module)
+    net_main_module.reset_flops_count = reset_flops_count.__get__(net_main_module)
+    net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__(
+                                                    net_main_module)
+
+    net_main_module.reset_flops_count()
+
+    return net_main_module
+
+def compute_average_flops_cost(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Returns current mean flops consumption per image.
+    """
+
+    for m in self.modules():
+        m.accumulate_flops = accumulate_flops.__get__(m)
+
+    layer_flops = {}
+    flops_sum = self.accumulate_flops(layer_flops)
+
+    for m in self.modules():
+        if hasattr(m, 'accumulate_flops'):
+            del m.accumulate_flops
+
+    layer_params = {}
+    for m in self.modules():
+        layer_params[m] = get_model_parameters_number(m)
+
+    params_sum = get_model_parameters_number(self)
+    return flops_sum / self.__batch_counter__, params_sum, layer_flops, layer_params
+
+
+def start_flops_count(self, **kwargs):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Activates the computation of mean flops consumption per image.
+    Call it before you run the network.
+    """
+    add_batch_counter_hook_function(self)
+
+    seen_types = set()
+
+    def add_flops_counter_hook_function(module, ost, verbose, ignore_list):
+        if type(module) in ignore_list:
+            seen_types.add(type(module))
+            if is_supported_instance(module):
+                module.__params__ = 0
+        elif is_supported_instance(module):
+            if hasattr(module, '__flops_handle__'):
+                return
+            if type(module) in CUSTOM_MODULES_MAPPING:
+                handle = module.register_forward_hook(
+                                        CUSTOM_MODULES_MAPPING[type(module)])
+            else:
+                handle = module.register_forward_hook(MODULES_MAPPING[type(module)])
+            module.__flops_handle__ = handle
+            seen_types.add(type(module))
+        else:
+            if verbose and not type(module) in (nn.Sequential, nn.ModuleList) and \
+               not type(module) in seen_types:
+                print('Warning: module ' + type(module).__name__ +
+                      ' is treated as a zero-op.', file=ost)
+            seen_types.add(type(module))
+
+    self.apply(partial(add_flops_counter_hook_function, **kwargs))
+
+
+def stop_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Stops computing the mean flops consumption per image.
+    Call whenever you want to pause the computation.
+    """
+    remove_batch_counter_hook_function(self)
+    self.apply(remove_flops_counter_hook_function)
+    self.apply(remove_flops_counter_variables)
+
+
+def reset_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Resets statistics computed so far.
+    """
+    add_batch_counter_variables_or_reset(self)
+    self.apply(add_flops_counter_variable_or_reset)
+
+
+# ---- Internal functions
+def batch_counter_hook(module, input, output):
+    batch_size = 1
+    if len(input) > 0:
+        # Can have multiple inputs, getting the first one
+        input = input[0]
+        batch_size = len(input)
+    else:
+        pass
+        print('Warning! No positional inputs found for a module,'
+              ' assuming batch size is 1.')
+    module.__batch_counter__ += batch_size
+
+
+def add_batch_counter_variables_or_reset(module):
+
+    module.__batch_counter__ = 0
+
+
+def add_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        return
+
+    handle = module.register_forward_hook(batch_counter_hook)
+    module.__batch_counter_handle__ = handle
+
+
+def remove_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        module.__batch_counter_handle__.remove()
+        del module.__batch_counter_handle__
+
+
+def add_flops_counter_variable_or_reset(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops__') or hasattr(module, '__params__'):
+            print('Warning: variables __flops__ or __params__ are already '
+                  'defined for the module' + type(module).__name__ +
+                  ' ptflops can affect your code!')
+            module.__ptflops_backup_flops__ = module.__flops__
+            module.__ptflops_backup_params__ = module.__params__
+        module.__flops__ = 0
+        module.__params__ = get_model_parameters_number(module)
+
+
+def is_supported_instance(module):
+    if type(module) in MODULES_MAPPING or type(module) in CUSTOM_MODULES_MAPPING:
+        return True
+    return False
+
+
+def remove_flops_counter_hook_function(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops_handle__'):
+            module.__flops_handle__.remove()
+            del module.__flops_handle__
+
+
+def remove_flops_counter_variables(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops__'):
+            del module.__flops__
+            if hasattr(module, '__ptflops_backup_flops__'):
+                module.__flops__ = module.__ptflops_backup_flops__
+        if hasattr(module, '__params__'):
+            del module.__params__
+            if hasattr(module, '__ptflops_backup_params__'):
+                module.__params__ = module.__ptflops_backup_params__
+    

ixformer_sdk/contrib/DeepCache/sd/pipeline_stable_diffusion.py

@@ -0,0 +1,812 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     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.
+import time 
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+import numpy as np
+from packaging import version
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    deprecate,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+
+from .unet_2d_condition import UNet2DConditionModel
+from .pipeline_utils import DiffusionPipeline
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import StableDiffusionPipeline
+
+        >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+        >>> pipe = pipe.to("cuda")
+
+        >>> prompt = "a photo of an astronaut riding a horse on mars"
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100):
+    samples = []
+    
+    while len(samples) < sample_size:
+        # Sample from a Gaussian centered at n/2
+        sample = int(np.random.normal(loc=n/2, scale=std_dev))
+        
+        # Check if the sample is in bounds
+        if 1 <= sample < n and sample not in samples:
+            samples.append(sample)
+    
+    return samples
+
+def sample_from_quad(total_numbers, n_samples, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1)
+
+        # Raise these values to the power of 1.5 to get a non-linear distribution
+        indices = np.unique(np.int32(x_values**pow))[:-1]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1)
+        indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+class StableDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin):
+    r"""
+    Pipeline for text-to-image generation using Stable Diffusion.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    The pipeline also inherits the following loading methods:
+        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
+        - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
+        - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
+        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
+        text_encoder ([`~transformers.CLIPTextModel`]):
+            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
+        tokenizer ([`~transformers.CLIPTokenizer`]):
+            A `CLIPTokenizer` to tokenize text.
+        unet ([`UNet2DConditionModel`]):
+            A `UNet2DConditionModel` to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offensive or harmful.
+            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
+            about a model's potential harms.
+        feature_extractor ([`~transformers.CLIPImageProcessor`]):
+            A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
+    """
+    model_cpu_offload_seq = "text_encoder->unet->vae"
+    _optional_components = ["safety_checker", "feature_extractor"]
+    _exclude_from_cpu_offload = ["safety_checker"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if safety_checker is None and requires_safety_checker:
+            logger.warning(
+                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
+                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
+                " results in services or applications open to the public. Both the diffusers team and Hugging Face"
+                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
+                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
+                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
+            )
+
+        if safety_checker is not None and feature_extractor is None:
+            raise ValueError(
+                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
+                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
+            )
+
+        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
+            version.parse(unet.config._diffusers_version).base_version
+        ) < version.parse("0.9.0.dev0")
+        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
+        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
+            deprecation_message = (
+                "The configuration file of the unet has set the default `sample_size` to smaller than"
+                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
+                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
+                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
+                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
+                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
+                " in the config might lead to incorrect results in future versions. If you have downloaded this"
+                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
+                " the `unet/config.json` file"
+            )
+            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(unet.config)
+            new_config["sample_size"] = 64
+            unet._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
+        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
+
+        prompt_embeds_tuple = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=lora_scale,
+        )
+
+        # concatenate for backwards comp
+        prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
+
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+                )
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            prompt_embeds_dtype = self.text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds, negative_prompt_embeds
+
+    def run_safety_checker(self, image, device, dtype):
+        if self.safety_checker is None:
+            has_nsfw_concept = None
+        else:
+            if torch.is_tensor(image):
+                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
+            else:
+                feature_extractor_input = self.image_processor.numpy_to_pil(image)
+            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
+            )
+        return image, has_nsfw_concept
+
+    def decode_latents(self, latents):
+        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
+        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents, return_dict=False)[0]
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        callback_steps,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        cache_interval: int = 1,
+        cache_layer_id: int = None,
+        cache_block_id: int = None,
+        uniform: bool = True,
+        pow: float = None,
+        center: int = None,
+        output_all_sequence: bool = False,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.7):
+                Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
+                using zero terminal SNR.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is a list with the generated images and the
+                second element is a list of `bool`s indicating whether the corresponding generated image contains
+                "not-safe-for-work" (nsfw) content.
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        prv_features = None
+        latents_list = [latents]
+
+        if cache_interval == 1:
+            interval_seq = list(range(num_inference_steps))
+        else:
+            if uniform:
+                interval_seq = list(range(0, num_inference_steps, cache_interval))
+            else:
+                num_slow_step = num_inference_steps//cache_interval
+                if num_inference_steps%cache_interval != 0:
+                    num_slow_step += 1
+                
+                interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+                #interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+        
+        interval_seq = sorted(interval_seq)
+        #print(interval_seq, len(interval_seq), pow)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            #print("[INFO] Update Feature Interval = {}, Update Layer Number = {}, Update Block Number = {}".format(cache_interval, cache_layer_id, cache_block_id))
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                if i in interval_seq:
+                    prv_features = None            
+
+                # predict the noise residual
+                noise_pred, prv_features = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    replicate_prv_feature=prv_features,
+                    quick_replicate= cache_interval>1,
+                    cache_layer_id=cache_layer_id,
+                    cache_block_id=cache_block_id,
+                    return_dict=False,
+                )
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                latents_list.append(latents)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        if not output_type == "latent":
+            if output_all_sequence:
+                image = [self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] for latents in latents_list]
+                has_nsfw_concept = None #self.run_safety_checker(images[0], device, prompt_embeds.dtype)
+                num_img = len(image)
+            else:
+                image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+                has_nsfw_concept = None
+                num_img = image.shape[0]
+        else:
+            image = latents
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * num_img
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        if output_all_sequence:
+            image = [self.image_processor.postprocess(img, output_type=output_type, do_denormalize=do_denormalize) for img in image]
+        else:
+            image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (image, has_nsfw_concept,)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

ixformer_sdk/contrib/DeepCache/sd/pipeline_text_to_video_zero.py

@@ -0,0 +1,741 @@
+import copy
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+from torch.nn.functional import grid_sample
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers.models import AutoencoderKL
+from .unet_2d_condition import UNet2DConditionModel
+from .pipeline_stable_diffusion import StableDiffusionPipeline, StableDiffusionSafetyChecker
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100):
+    samples = []
+    
+    while len(samples) < sample_size:
+        # Sample from a Gaussian centered at n/2
+        sample = int(np.random.normal(loc=n/2, scale=std_dev))
+        
+        # Check if the sample is in bounds
+        if 1 <= sample < n and sample not in samples:
+            samples.append(sample)
+    
+    return samples
+
+def sample_from_quad(total_numbers, n_samples, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1)
+
+        # Raise these values to the power of 1.5 to get a non-linear distribution
+        indices = np.unique(np.int32(x_values**pow))[:-1]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1)
+        indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def rearrange_0(tensor, f):
+    F, C, H, W = tensor.size()
+    tensor = torch.permute(torch.reshape(tensor, (F // f, f, C, H, W)), (0, 2, 1, 3, 4))
+    return tensor
+
+
+def rearrange_1(tensor):
+    B, C, F, H, W = tensor.size()
+    return torch.reshape(torch.permute(tensor, (0, 2, 1, 3, 4)), (B * F, C, H, W))
+
+
+def rearrange_3(tensor, f):
+    F, D, C = tensor.size()
+    return torch.reshape(tensor, (F // f, f, D, C))
+
+
+def rearrange_4(tensor):
+    B, F, D, C = tensor.size()
+    return torch.reshape(tensor, (B * F, D, C))
+
+
+class CrossFrameAttnProcessor:
+    """
+    Cross frame attention processor. Each frame attends the first frame.
+
+    Args:
+        batch_size: The number that represents actual batch size, other than the frames.
+            For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to
+            2, due to classifier-free guidance.
+    """
+
+    def __init__(self, batch_size=2):
+        self.batch_size = batch_size
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+        query = attn.to_q(hidden_states)
+
+        is_cross_attention = encoder_hidden_states is not None
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        # Cross Frame Attention
+        if not is_cross_attention:
+            video_length = key.size()[0] // self.batch_size
+            first_frame_index = [0] * video_length
+
+            # rearrange keys to have batch and frames in the 1st and 2nd dims respectively
+            key = rearrange_3(key, video_length)
+            key = key[:, first_frame_index]
+            # rearrange values to have batch and frames in the 1st and 2nd dims respectively
+            value = rearrange_3(value, video_length)
+            value = value[:, first_frame_index]
+
+            # rearrange back to original shape
+            key = rearrange_4(key)
+            value = rearrange_4(value)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class CrossFrameAttnProcessor2_0:
+    """
+    Cross frame attention processor with scaled_dot_product attention of Pytorch 2.0.
+
+    Args:
+        batch_size: The number that represents actual batch size, other than the frames.
+            For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to
+            2, due to classifier-free guidance.
+    """
+
+    def __init__(self, batch_size=2):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+        self.batch_size = batch_size
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        inner_dim = hidden_states.shape[-1]
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        query = attn.to_q(hidden_states)
+
+        is_cross_attention = encoder_hidden_states is not None
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        # Cross Frame Attention
+        if not is_cross_attention:
+            video_length = max(1, key.size()[0] // self.batch_size)
+            first_frame_index = [0] * video_length
+
+            # rearrange keys to have batch and frames in the 1st and 2nd dims respectively
+            key = rearrange_3(key, video_length)
+            key = key[:, first_frame_index]
+            # rearrange values to have batch and frames in the 1st and 2nd dims respectively
+            value = rearrange_3(value, video_length)
+            value = value[:, first_frame_index]
+
+            # rearrange back to original shape
+            key = rearrange_4(key)
+            value = rearrange_4(value)
+
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+        return hidden_states
+
+
+@dataclass
+class TextToVideoPipelineOutput(BaseOutput):
+    r"""
+    Output class for zero-shot text-to-video pipeline.
+
+    Args:
+        images (`[List[PIL.Image.Image]`, `np.ndarray`]):
+            List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+            num_channels)`.
+        nsfw_content_detected (`[List[bool]]`):
+            List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
+            `None` if safety checking could not be performed.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+    nsfw_content_detected: Optional[List[bool]]
+
+
+def coords_grid(batch, ht, wd, device):
+    # Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py
+    coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
+    coords = torch.stack(coords[::-1], dim=0).float()
+    return coords[None].repeat(batch, 1, 1, 1)
+
+
+def warp_single_latent(latent, reference_flow):
+    """
+    Warp latent of a single frame with given flow
+
+    Args:
+        latent: latent code of a single frame
+        reference_flow: flow which to warp the latent with
+
+    Returns:
+        warped: warped latent
+    """
+    _, _, H, W = reference_flow.size()
+    _, _, h, w = latent.size()
+    coords0 = coords_grid(1, H, W, device=latent.device).to(latent.dtype)
+
+    coords_t0 = coords0 + reference_flow
+    coords_t0[:, 0] /= W
+    coords_t0[:, 1] /= H
+
+    coords_t0 = coords_t0 * 2.0 - 1.0
+    coords_t0 = F.interpolate(coords_t0, size=(h, w), mode="bilinear")
+    coords_t0 = torch.permute(coords_t0, (0, 2, 3, 1))
+
+    warped = grid_sample(latent, coords_t0, mode="nearest", padding_mode="reflection")
+    return warped
+
+
+def create_motion_field(motion_field_strength_x, motion_field_strength_y, frame_ids, device, dtype):
+    """
+    Create translation motion field
+
+    Args:
+        motion_field_strength_x: motion strength along x-axis
+        motion_field_strength_y: motion strength along y-axis
+        frame_ids: indexes of the frames the latents of which are being processed.
+            This is needed when we perform chunk-by-chunk inference
+        device: device
+        dtype: dtype
+
+    Returns:
+
+    """
+    seq_length = len(frame_ids)
+    reference_flow = torch.zeros((seq_length, 2, 512, 512), device=device, dtype=dtype)
+    for fr_idx in range(seq_length):
+        reference_flow[fr_idx, 0, :, :] = motion_field_strength_x * (frame_ids[fr_idx])
+        reference_flow[fr_idx, 1, :, :] = motion_field_strength_y * (frame_ids[fr_idx])
+    return reference_flow
+
+
+def create_motion_field_and_warp_latents(motion_field_strength_x, motion_field_strength_y, frame_ids, latents):
+    """
+    Creates translation motion and warps the latents accordingly
+
+    Args:
+        motion_field_strength_x: motion strength along x-axis
+        motion_field_strength_y: motion strength along y-axis
+        frame_ids: indexes of the frames the latents of which are being processed.
+            This is needed when we perform chunk-by-chunk inference
+        latents: latent codes of frames
+
+    Returns:
+        warped_latents: warped latents
+    """
+    motion_field = create_motion_field(
+        motion_field_strength_x=motion_field_strength_x,
+        motion_field_strength_y=motion_field_strength_y,
+        frame_ids=frame_ids,
+        device=latents.device,
+        dtype=latents.dtype,
+    )
+    warped_latents = latents.clone().detach()
+    for i in range(len(warped_latents)):
+        warped_latents[i] = warp_single_latent(latents[i][None], motion_field[i][None])
+    return warped_latents
+
+
+class TextToVideoZeroPipeline(StableDiffusionPipeline):
+    r"""
+    Pipeline for zero-shot text-to-video generation using Stable Diffusion.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
+        tokenizer (`CLIPTokenizer`):
+            A [`~transformers.CLIPTokenizer`] to tokenize text.
+        unet ([`UNet2DConditionModel`]):
+            A [`UNet3DConditionModel`] to denoise the encoded video latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offensive or harmful.
+            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
+            about a model's potential harms.
+        feature_extractor ([`CLIPImageProcessor`]):
+            A [`CLIPImageProcessor`] to extract features from generated images; used as inputs to the `safety_checker`.
+    """
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__(
+            vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker
+        )
+        processor = (
+            CrossFrameAttnProcessor2_0(batch_size=2)
+            if hasattr(F, "scaled_dot_product_attention")
+            else CrossFrameAttnProcessor(batch_size=2)
+        )
+        self.unet.set_attn_processor(processor)
+
+    def forward_loop(self, x_t0, t0, t1, generator):
+        """
+        Perform DDPM forward process from time t0 to t1. This is the same as adding noise with corresponding variance.
+
+        Args:
+            x_t0:
+                Latent code at time t0.
+            t0:
+                Timestep at t0.
+            t1:
+                Timestamp at t1.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+
+        Returns:
+            x_t1:
+                Forward process applied to x_t0 from time t0 to t1.
+        """
+        eps = randn_tensor(x_t0.size(), generator=generator, dtype=x_t0.dtype, device=x_t0.device)
+        alpha_vec = torch.prod(self.scheduler.alphas[t0:t1])
+        x_t1 = torch.sqrt(alpha_vec) * x_t0 + torch.sqrt(1 - alpha_vec) * eps
+        return x_t1
+
+    def backward_loop(
+        self,
+        latents,
+        timesteps,
+        prompt_embeds,
+        guidance_scale,
+        callback,
+        callback_steps,
+        num_warmup_steps,
+        extra_step_kwargs,
+        prv_features,
+        interval_seq,
+        cache_interval,
+        cache_block_id,
+        cache_layer_id,
+        cross_attention_kwargs=None,
+    ):
+        """
+        Perform backward process given list of time steps.
+
+        Args:
+            latents:
+                Latents at time timesteps[0].
+            timesteps:
+                Time steps along which to perform backward process.
+            prompt_embeds:
+                Pre-generated text embeddings.
+            guidance_scale:
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            extra_step_kwargs:
+                Extra_step_kwargs.
+            cross_attention_kwargs:
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            num_warmup_steps:
+                number of warmup steps.
+
+        Returns:
+            latents:
+                Latents of backward process output at time timesteps[-1].
+        """
+        do_classifier_free_guidance = guidance_scale > 1.0
+        num_steps = (len(timesteps) - num_warmup_steps) // self.scheduler.order
+        with self.progress_bar(total=num_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                ########
+                if i in interval_seq:
+                    prv_features = None
+                # predict the noise residual
+                noise_pred, prv_features = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    replicate_prv_feature=prv_features,
+                    quick_replicate= cache_interval>1,
+                    cache_layer_id=cache_layer_id,
+                    cache_block_id=cache_block_id,
+                    return_dict=False,
+                )
+                ########
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+        return latents.clone().detach()
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        video_length: Optional[int] = 8,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        motion_field_strength_x: float = 12,
+        motion_field_strength_y: float = 12,
+        output_type: Optional[str] = "tensor",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+        t0: int = 44,
+        t1: int = 47,
+        frame_ids: Optional[List[int]] = None,
+        ########
+        cache_interval: int = 1,
+        cache_layer_id: int = None,
+        cache_block_id: int = None,
+        uniform: bool = True,
+        pow: float = None,
+        center: int = None,
+        output_all_sequence: bool = False,
+        ########
+    ):
+        """
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            video_length (`int`, *optional*, defaults to 8):
+                The number of generated video frames.
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in video generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of videos to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for video
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            output_type (`str`, *optional*, defaults to `"numpy"`):
+                The output format of the generated video. Choose between `"latent"` and `"numpy"`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a
+                [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`] instead of
+                a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            motion_field_strength_x (`float`, *optional*, defaults to 12):
+                Strength of motion in generated video along x-axis. See the [paper](https://arxiv.org/abs/2303.13439),
+                Sect. 3.3.1.
+            motion_field_strength_y (`float`, *optional*, defaults to 12):
+                Strength of motion in generated video along y-axis. See the [paper](https://arxiv.org/abs/2303.13439),
+                Sect. 3.3.1.
+            t0 (`int`, *optional*, defaults to 44):
+                Timestep t0. Should be in the range [0, num_inference_steps - 1]. See the
+                [paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1.
+            t1 (`int`, *optional*, defaults to 47):
+                Timestep t0. Should be in the range [t0 + 1, num_inference_steps - 1]. See the
+                [paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1.
+            frame_ids (`List[int]`, *optional*):
+                Indexes of the frames that are being generated. This is used when generating longer videos
+                chunk-by-chunk.
+
+        Returns:
+            [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`]:
+                The output contains a `ndarray` of the generated video, when `output_type` != `"latent"`, otherwise a
+                latent code of generated videos and a list of `bool`s indicating whether the corresponding generated
+                video contains "not-safe-for-work" (nsfw) content..
+        """
+        assert video_length > 0
+        if frame_ids is None:
+            frame_ids = list(range(video_length))
+        assert len(frame_ids) == video_length
+
+        assert num_videos_per_prompt == 1
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if isinstance(negative_prompt, str):
+            negative_prompt = [negative_prompt]
+
+        # Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # Check inputs. Raise error if not correct
+        self.check_inputs(prompt, height, width, callback_steps)
+
+        # Define call parameters
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # Encode input prompt
+        prompt_embeds = self._encode_prompt(
+            prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
+        )
+
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        # Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        prv_features = None #record cache feature ****
+        latents_list = [latents] 
+
+        if cache_interval == 1:
+            interval_seq = list(range(num_inference_steps))
+        else:
+            if uniform:
+                interval_seq = list(range(0, num_inference_steps, cache_interval))
+            else:
+                num_slow_step = num_inference_steps//cache_interval
+                if num_inference_steps%cache_interval != 0:
+                    num_slow_step += 1
+                
+                interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+                #interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+        
+        interval_seq = sorted(interval_seq)
+
+        # Perform the first backward process up to time T_1
+        x_1_t1 = self.backward_loop(
+            timesteps=timesteps[: -t1 - 1],
+            prompt_embeds=prompt_embeds,
+            latents=latents,
+            guidance_scale=guidance_scale,
+            callback=callback,
+            callback_steps=callback_steps,
+            extra_step_kwargs=extra_step_kwargs,
+            num_warmup_steps=num_warmup_steps,
+            prv_features=prv_features,
+            interval_seq=interval_seq,
+            cache_interval=cache_interval,
+            cache_block_id=cache_block_id,
+            cache_layer_id=cache_layer_id,
+        )
+        scheduler_copy = copy.deepcopy(self.scheduler)
+
+        # Perform the second backward process up to time T_0
+        x_1_t0 = self.backward_loop(
+            timesteps=timesteps[-t1 - 1 : -t0 - 1],
+            prompt_embeds=prompt_embeds,
+            latents=x_1_t1,
+            guidance_scale=guidance_scale,
+            callback=callback,
+            callback_steps=callback_steps,
+            extra_step_kwargs=extra_step_kwargs,
+            num_warmup_steps=0,
+            prv_features=prv_features,
+            interval_seq=interval_seq,
+            cache_interval=cache_interval,
+            cache_block_id=cache_block_id,
+            cache_layer_id=cache_layer_id,
+        )
+
+        # Propagate first frame latents at time T_0 to remaining frames
+        x_2k_t0 = x_1_t0.repeat(video_length - 1, 1, 1, 1)
+
+        # Add motion in latents at time T_0
+        x_2k_t0 = create_motion_field_and_warp_latents(
+            motion_field_strength_x=motion_field_strength_x,
+            motion_field_strength_y=motion_field_strength_y,
+            latents=x_2k_t0,
+            frame_ids=frame_ids[1:],
+        )
+
+        # Perform forward process up to time T_1
+        x_2k_t1 = self.forward_loop(
+            x_t0=x_2k_t0,
+            t0=timesteps[-t0 - 1].item(),
+            t1=timesteps[-t1 - 1].item(),
+            generator=generator,
+        )
+
+        # Perform backward process from time T_1 to 0
+        x_1k_t1 = torch.cat([x_1_t1, x_2k_t1])
+        b, l, d = prompt_embeds.size()
+        prompt_embeds = prompt_embeds[:, None].repeat(1, video_length, 1, 1).reshape(b * video_length, l, d)
+
+        self.scheduler = scheduler_copy
+        x_1k_0 = self.backward_loop(
+            timesteps=timesteps[-t1 - 1 :],
+            prompt_embeds=prompt_embeds,
+            latents=x_1k_t1,
+            guidance_scale=guidance_scale,
+            callback=callback,
+            callback_steps=callback_steps,
+            extra_step_kwargs=extra_step_kwargs,
+            num_warmup_steps=0,
+            prv_features=prv_features,
+            interval_seq=interval_seq,
+            cache_interval=cache_interval,
+            cache_block_id=cache_block_id,
+            cache_layer_id=cache_layer_id,
+        )
+        latents = x_1k_0
+
+        # manually for max memory savings
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.unet.to("cpu")
+        torch.cuda.empty_cache()
+
+        if output_type == "latent":
+            image = latents
+            has_nsfw_concept = None
+        else:
+            image = self.decode_latents(latents)
+            # Run safety checker
+            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return TextToVideoPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

ixformer_sdk/contrib/DeepCache/svd/pipeline_stable_video_diffusion.py

@@ -0,0 +1,659 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     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.
+
+import inspect
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
+from diffusers.schedulers import EulerDiscreteScheduler
+from diffusers.utils import BaseOutput, logging
+from diffusers.utils.torch_utils import randn_tensor
+from .pipeline_utils import DiffusionPipeline
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def _append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
+    return x[(...,) + (None,) * dims_to_append]
+
+
+def tensor2vid(video: torch.Tensor, processor, output_type="np"):
+    # Based on:
+    # https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78
+
+    batch_size, channels, num_frames, height, width = video.shape
+    outputs = []
+    for batch_idx in range(batch_size):
+        batch_vid = video[batch_idx].permute(1, 0, 2, 3)
+        batch_output = processor.postprocess(batch_vid, output_type)
+
+        outputs.append(batch_output)
+
+    return outputs
+
+
+@dataclass
+class StableVideoDiffusionPipelineOutput(BaseOutput):
+    r"""
+    Output class for zero-shot text-to-video pipeline.
+
+    Args:
+        frames (`[List[PIL.Image.Image]`, `np.ndarray`]):
+            List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+            num_channels)`.
+    """
+
+    frames: Union[List[PIL.Image.Image], np.ndarray]
+
+
+class StableVideoDiffusionPipeline(DiffusionPipeline):
+    r"""
+    Pipeline to generate video from an input image using Stable Video Diffusion.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
+        image_encoder ([`~transformers.CLIPVisionModelWithProjection`]):
+            Frozen CLIP image-encoder ([laion/CLIP-ViT-H-14-laion2B-s32B-b79K](https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K)).
+        unet ([`UNetSpatioTemporalConditionModel`]):cache_interval=5, cache_branch=0,
+            A `UNetSpatioTemporalConditionModel` to denoise the encoded image latents.
+        scheduler ([`EulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+        feature_extractor ([`~transformers.CLIPImageProcessor`]):
+            A `CLIPImageProcessor` to extract features from generated images.
+    """
+
+    model_cpu_offload_seq = "image_encoder->unet->vae"
+    _callback_tensor_inputs = ["latents"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKLTemporalDecoder,
+        image_encoder: CLIPVisionModelWithProjection,
+        unet: UNetSpatioTemporalConditionModel,
+        scheduler: EulerDiscreteScheduler,
+        feature_extractor: CLIPImageProcessor,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            image_encoder=image_encoder,
+            unet=unet,
+            scheduler=scheduler,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    def _encode_image(self, image, device, num_videos_per_prompt, do_classifier_free_guidance):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if not isinstance(image, torch.Tensor):
+            image = self.image_processor.pil_to_numpy(image)
+            image = self.image_processor.numpy_to_pt(image)
+
+            # We normalize the image before resizing to match with the original implementation.
+            # Then we unnormalize it after resizing.
+            image = image * 2.0 - 1.0
+            image = _resize_with_antialiasing(image, (224, 224))
+            image = (image + 1.0) / 2.0
+
+            # Normalize the image with for CLIP input
+            image = self.feature_extractor(
+                images=image,
+                do_normalize=True,
+                do_center_crop=False,
+                do_resize=False,
+                do_rescale=False,
+                return_tensors="pt",
+            ).pixel_values
+
+        image = image.to(device=device, dtype=dtype)
+        image_embeddings = self.image_encoder(image).image_embeds
+        image_embeddings = image_embeddings.unsqueeze(1)
+
+        # duplicate image embeddings for each generation per prompt, using mps friendly method
+        bs_embed, seq_len, _ = image_embeddings.shape
+        image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1)
+        image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
+
+        if do_classifier_free_guidance:
+            negative_image_embeddings = torch.zeros_like(image_embeddings)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            image_embeddings = torch.cat([negative_image_embeddings, image_embeddings])
+
+        return image_embeddings
+
+    def _encode_vae_image(
+        self,
+        image: torch.Tensor,
+        device,
+        num_videos_per_prompt,
+        do_classifier_free_guidance,
+    ):
+        image = image.to(device=device)
+        image_latents = self.vae.encode(image).latent_dist.mode()
+
+        if do_classifier_free_guidance:
+            negative_image_latents = torch.zeros_like(image_latents)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            image_latents = torch.cat([negative_image_latents, image_latents])
+
+        # duplicate image_latents for each generation per prompt, using mps friendly method
+        image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1)
+
+        return image_latents
+
+    def _get_add_time_ids(
+        self,
+        fps,
+        motion_bucket_id,
+        noise_aug_strength,
+        dtype,
+        batch_size,
+        num_videos_per_prompt,
+        do_classifier_free_guidance,
+    ):
+        add_time_ids = [fps, motion_bucket_id, noise_aug_strength]
+
+        passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids)
+        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
+
+        if expected_add_embed_dim != passed_add_embed_dim:
+            raise ValueError(
+                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+            )
+
+        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+        add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1)
+
+        if do_classifier_free_guidance:
+            add_time_ids = torch.cat([add_time_ids, add_time_ids])
+
+        return add_time_ids
+
+    def decode_latents(self, latents, num_frames, decode_chunk_size=14):
+        # [batch, frames, channels, height, width] -> [batch*frames, channels, height, width]
+        latents = latents.flatten(0, 1)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+
+        accepts_num_frames = "num_frames" in set(inspect.signature(self.vae.forward).parameters.keys())
+
+        # decode decode_chunk_size frames at a time to avoid OOM
+        frames = []
+        for i in range(0, latents.shape[0], decode_chunk_size):
+            num_frames_in = latents[i : i + decode_chunk_size].shape[0]
+            decode_kwargs = {}
+            if accepts_num_frames:
+                # we only pass num_frames_in if it's expected
+                decode_kwargs["num_frames"] = num_frames_in
+
+            frame = self.vae.decode(latents[i : i + decode_chunk_size], **decode_kwargs).sample
+            frames.append(frame)
+        frames = torch.cat(frames, dim=0)
+
+        # [batch*frames, channels, height, width] -> [batch, channels, frames, height, width]
+        frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)
+
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        frames = frames.float()
+        return frames
+
+    def check_inputs(self, image, height, width):
+        if (
+            not isinstance(image, torch.Tensor)
+            and not isinstance(image, PIL.Image.Image)
+            and not isinstance(image, list)
+        ):
+            raise ValueError(
+                "`image` has to be of type `torch.FloatTensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+                f" {type(image)}"
+            )
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_frames,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        shape = (
+            batch_size,
+            num_frames,
+            num_channels_latents // 2,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.FloatTensor],
+        height: int = 576,
+        width: int = 1024,
+        num_frames: Optional[int] = None,
+        num_inference_steps: int = 25,
+        min_guidance_scale: float = 1.0,
+        max_guidance_scale: float = 3.0,
+        fps: int = 7,
+        motion_bucket_id: int = 127,
+        noise_aug_strength: int = 0.02,
+        decode_chunk_size: Optional[int] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        cache_interval: Optional[int] = 1, 
+        cache_branch: Optional[int] = None,
+        return_dict: bool = True,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
+                Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
+                [`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_frames (`int`, *optional*):
+                The number of video frames to generate. Defaults to 14 for `stable-video-diffusion-img2vid` and to 25 for `stable-video-diffusion-img2vid-xt`
+            num_inference_steps (`int`, *optional*, defaults to 25):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference. This parameter is modulated by `strength`.
+            min_guidance_scale (`float`, *optional*, defaults to 1.0):
+                The minimum guidance scale. Used for the classifier free guidance with first frame.
+            max_guidance_scale (`float`, *optional*, defaults to 3.0):
+                The maximum guidance scale. Used for the classifier free guidance with last frame.
+            fps (`int`, *optional*, defaults to 7):
+                Frames per second. The rate at which the generated images shall be exported to a video after generation.
+                Note that Stable Diffusion Video's UNet was micro-conditioned on fps-1 during training.
+            motion_bucket_id (`int`, *optional*, defaults to 127):
+                The motion bucket ID. Used as conditioning for the generation. The higher the number the more motion will be in the video.
+            noise_aug_strength (`int`, *optional*, defaults to 0.02):
+                The amount of noise added to the init image, the higher it is the less the video will look like the init image. Increase it for more motion.
+            decode_chunk_size (`int`, *optional*):
+                The number of frames to decode at a time. The higher the chunk size, the higher the temporal consistency
+                between frames, but also the higher the memory consumption. By default, the decoder will decode all frames at once
+                for maximal quality. Reduce `decode_chunk_size` to reduce memory usage.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is a list of list with the generated frames.
+
+        Examples:
+
+        ```py
+        from diffusers import StableVideoDiffusionPipeline
+        from diffusers.utils import load_image, export_to_video
+
+        pipe = StableVideoDiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16")
+        pipe.to("cuda")
+
+        image = load_image("https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200")
+        image = image.resize((1024, 576))
+
+        frames = pipe(image, num_frames=25, decode_chunk_size=8).frames[0]
+        export_to_video(frames, "generated.mp4", fps=7)
+        ```
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        num_frames = num_frames if num_frames is not None else self.unet.config.num_frames
+        decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(image, height, width)
+
+        # 2. Define call parameters
+        if isinstance(image, PIL.Image.Image):
+            batch_size = 1
+        elif isinstance(image, list):
+            batch_size = len(image)
+        else:
+            batch_size = image.shape[0]
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = max_guidance_scale > 1.0
+
+        # 3. Encode input image
+        image_embeddings = self._encode_image(image, device, num_videos_per_prompt, do_classifier_free_guidance)
+
+        # NOTE: Stable Diffusion Video was conditioned on fps - 1, which
+        # is why it is reduced here.
+        # See: https://github.com/Stability-AI/generative-models/blob/ed0997173f98eaf8f4edf7ba5fe8f15c6b877fd3/scripts/sampling/simple_video_sample.py#L188
+        fps = fps - 1
+
+        # 4. Encode input image using VAE
+        image = self.image_processor.preprocess(image, height=height, width=width)
+        noise = randn_tensor(image.shape, generator=generator, device=image.device, dtype=image.dtype)
+        image = image + noise_aug_strength * noise
+
+        needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float32)
+
+        image_latents = self._encode_vae_image(image, device, num_videos_per_prompt, do_classifier_free_guidance)
+        image_latents = image_latents.to(image_embeddings.dtype)
+
+        # cast back to fp16 if needed
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float16)
+
+        # Repeat the image latents for each frame so we can concatenate them with the noise
+        # image_latents [batch, channels, height, width] ->[batch, num_frames, channels, height, width]
+        image_latents = image_latents.unsqueeze(1).repeat(1, num_frames, 1, 1, 1)
+
+        # 5. Get Added Time IDs
+        added_time_ids = self._get_add_time_ids(
+            fps,
+            motion_bucket_id,
+            noise_aug_strength,
+            image_embeddings.dtype,
+            batch_size,
+            num_videos_per_prompt,
+            do_classifier_free_guidance,
+        )
+        added_time_ids = added_time_ids.to(device)
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_frames,
+            num_channels_latents,
+            height,
+            width,
+            image_embeddings.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 7. Prepare guidance scale
+        guidance_scale = torch.linspace(min_guidance_scale, max_guidance_scale, num_frames).unsqueeze(0)
+        guidance_scale = guidance_scale.to(device, latents.dtype)
+        guidance_scale = guidance_scale.repeat(batch_size * num_videos_per_prompt, 1)
+        guidance_scale = _append_dims(guidance_scale, latents.ndim)
+
+        self._guidance_scale = guidance_scale
+
+        cache_features = None
+        interval_seq = list(range(0, num_inference_steps, cache_interval))
+        interval_seq = sorted(interval_seq)
+
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # Concatenate image_latents over channels dimention
+                latent_model_input = torch.cat([latent_model_input, image_latents], dim=2)
+
+                if i in interval_seq:
+                    cache_features = None
+
+                # predict the noise residual
+                noise_pred, cache_features = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=image_embeddings,
+                    added_time_ids=added_time_ids,
+                    cache_features=cache_features,
+                    cache_branch=cache_branch,
+                    return_dict=False,
+                )
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        if not output_type == "latent":
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                self.vae.to(dtype=torch.float16)
+            frames = self.decode_latents(latents, num_frames, decode_chunk_size)
+            frames = tensor2vid(frames, self.image_processor, output_type=output_type)
+        else:
+            frames = latents
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return frames
+
+        return StableVideoDiffusionPipelineOutput(frames=frames)
+
+
+# resizing utils
+# TODO: clean up later
+def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True):
+    h, w = input.shape[-2:]
+    factors = (h / size[0], w / size[1])
+
+    # First, we have to determine sigma
+    # Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171
+    sigmas = (
+        max((factors[0] - 1.0) / 2.0, 0.001),
+        max((factors[1] - 1.0) / 2.0, 0.001),
+    )
+
+    # Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma
+    # https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206
+    # But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now
+    ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3))
+
+    # Make sure it is odd
+    if (ks[0] % 2) == 0:
+        ks = ks[0] + 1, ks[1]
+
+    if (ks[1] % 2) == 0:
+        ks = ks[0], ks[1] + 1
+
+    input = _gaussian_blur2d(input, ks, sigmas)
+
+    output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
+    return output
+
+
+def _compute_padding(kernel_size):
+    """Compute padding tuple."""
+    # 4 or 6 ints:  (padding_left, padding_right,padding_top,padding_bottom)
+    # https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad
+    if len(kernel_size) < 2:
+        raise AssertionError(kernel_size)
+    computed = [k - 1 for k in kernel_size]
+
+    # for even kernels we need to do asymmetric padding :(
+    out_padding = 2 * len(kernel_size) * [0]
+
+    for i in range(len(kernel_size)):
+        computed_tmp = computed[-(i + 1)]
+
+        pad_front = computed_tmp // 2
+        pad_rear = computed_tmp - pad_front
+
+        out_padding[2 * i + 0] = pad_front
+        out_padding[2 * i + 1] = pad_rear
+
+    return out_padding
+
+
+def _filter2d(input, kernel):
+    # prepare kernel
+    b, c, h, w = input.shape
+    tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype)
+
+    tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)
+
+    height, width = tmp_kernel.shape[-2:]
+
+    padding_shape: list[int] = _compute_padding([height, width])
+    input = torch.nn.functional.pad(input, padding_shape, mode="reflect")
+
+    # kernel and input tensor reshape to align element-wise or batch-wise params
+    tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
+    input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1))
+
+    # convolve the tensor with the kernel.
+    output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
+
+    out = output.view(b, c, h, w)
+    return out
+
+
+def _gaussian(window_size: int, sigma):
+    if isinstance(sigma, float):
+        sigma = torch.tensor([[sigma]])
+
+    batch_size = sigma.shape[0]
+
+    x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1)
+
+    if window_size % 2 == 0:
+        x = x + 0.5
+
+    gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
+
+    return gauss / gauss.sum(-1, keepdim=True)
+
+
+def _gaussian_blur2d(input, kernel_size, sigma):
+    if isinstance(sigma, tuple):
+        sigma = torch.tensor([sigma], dtype=input.dtype)
+    else:
+        sigma = sigma.to(dtype=input.dtype)
+
+    ky, kx = int(kernel_size[0]), int(kernel_size[1])
+    bs = sigma.shape[0]
+    kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1))
+    kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1))
+    out_x = _filter2d(input, kernel_x[..., None, :])
+    out = _filter2d(out_x, kernel_y[..., None])
+
+    return out

ixformer_sdk/contrib/DeepCache/svd/unet_spatio_temporal_condition.py

@@ -0,0 +1,566 @@
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import UNet2DConditionLoadersMixin
+from diffusers.utils import BaseOutput, logging
+from diffusers.models.attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+
+from .unet_3d_blocks import UNetMidBlockSpatioTemporal, get_down_block, get_up_block
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNetSpatioTemporalConditionOutput(BaseOutput):
+    """
+    The output of [`UNetSpatioTemporalConditionModel`].
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):
+            The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor = None
+
+
+class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
+    r"""
+    A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and returns a sample
+    shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample.
+        in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
+            The tuple of downsample blocks to use.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
+            The tuple of upsample blocks to use.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        addition_time_embed_dim: (`int`, defaults to 256):
+            Dimension to to encode the additional time ids.
+        projection_class_embeddings_input_dim (`int`, defaults to 768):
+            The dimension of the projection of encoded `added_time_ids`.
+        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
+        cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+            The dimension of the cross attention features.
+        transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
+            [`~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`], [`~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`],
+            [`~models.unet_3d_blocks.UNetMidBlockSpatioTemporal`].
+        num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`):
+            The number of attention heads.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 8,
+        out_channels: int = 4,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlockSpatioTemporal",
+            "CrossAttnDownBlockSpatioTemporal",
+            "CrossAttnDownBlockSpatioTemporal",
+            "DownBlockSpatioTemporal",
+        ),
+        up_block_types: Tuple[str] = (
+            "UpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+        ),
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        addition_time_embed_dim: int = 256,
+        projection_class_embeddings_input_dim: int = 768,
+        layers_per_block: Union[int, Tuple[int]] = 2,
+        cross_attention_dim: Union[int, Tuple[int]] = 1024,
+        transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
+        num_attention_heads: Union[int, Tuple[int]] = (5, 10, 10, 20),
+        num_frames: int = 25,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+
+        # Check inputs
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+            )
+
+        if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
+            )
+
+        # input
+        self.conv_in = nn.Conv2d(
+            in_channels,
+            block_out_channels[0],
+            kernel_size=3,
+            padding=1,
+        )
+
+        # time
+        time_embed_dim = block_out_channels[0] * 4
+
+        self.time_proj = Timesteps(block_out_channels[0], True, downscale_freq_shift=0)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+
+        self.add_time_proj = Timesteps(addition_time_embed_dim, True, downscale_freq_shift=0)
+        self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        if isinstance(cross_attention_dim, int):
+            cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
+
+        if isinstance(layers_per_block, int):
+            layers_per_block = [layers_per_block] * len(down_block_types)
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        blocks_time_embed_dim = time_embed_dim
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block[i],
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=1e-5,
+                cross_attention_dim=cross_attention_dim[i],
+                num_attention_heads=num_attention_heads[i],
+                resnet_act_fn="silu",
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidBlockSpatioTemporal(
+            block_out_channels[-1],
+            temb_channels=blocks_time_embed_dim,
+            transformer_layers_per_block=transformer_layers_per_block[-1],
+            cross_attention_dim=cross_attention_dim[-1],
+            num_attention_heads=num_attention_heads[-1],
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_num_attention_heads = list(reversed(num_attention_heads))
+        reversed_layers_per_block = list(reversed(layers_per_block))
+        reversed_cross_attention_dim = list(reversed(cross_attention_dim))
+        reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
+
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=reversed_layers_per_block[i] + 1,
+                transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=1e-5,
+                resolution_idx=i,
+                cross_attention_dim=reversed_cross_attention_dim[i],
+                num_attention_heads=reversed_num_attention_heads[i],
+                resnet_act_fn="silu",
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-5)
+        self.conv_act = nn.SiLU()
+
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0],
+            out_channels,
+            kernel_size=3,
+            padding=1,
+        )
+
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    # Copied from diffusers.models.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
+    def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+        """
+        Sets the attention processor to use [feed forward
+        chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
+
+        Parameters:
+            chunk_size (`int`, *optional*):
+                The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
+                over each tensor of dim=`dim`.
+            dim (`int`, *optional*, defaults to `0`):
+                The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
+                or dim=1 (sequence length).
+        """
+        if dim not in [0, 1]:
+            raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
+
+        # By default chunk size is 1
+        chunk_size = chunk_size or 1
+
+        def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
+            if hasattr(module, "set_chunk_feed_forward"):
+                module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
+
+            for child in module.children():
+                fn_recursive_feed_forward(child, chunk_size, dim)
+
+        for module in self.children():
+            fn_recursive_feed_forward(module, chunk_size, dim)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        added_time_ids: torch.Tensor,
+        cache_features: Optional[torch.Tensor] = None,
+        cache_branch: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]:
+        r"""
+        The [`UNetSpatioTemporalConditionModel`] forward method.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`.
+            timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.FloatTensor`):
+                The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`.
+            added_time_ids: (`torch.FloatTensor`):
+                The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal
+                embeddings and added to the time embeddings.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead of a plain
+                tuple.
+        Returns:
+            [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`:
+                If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is returned, otherwise
+                a `tuple` is returned where the first element is the sample tensor.
+        """
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        batch_size, num_frames = sample.shape[:2]
+        timesteps = timesteps.expand(batch_size)
+
+        t_emb = self.time_proj(timesteps)
+
+        # `Timesteps` does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb)
+
+        time_embeds = self.add_time_proj(added_time_ids.flatten())
+        time_embeds = time_embeds.reshape((batch_size, -1))
+        time_embeds = time_embeds.to(emb.dtype)
+        aug_emb = self.add_embedding(time_embeds)
+        emb = emb + aug_emb
+
+        # Flatten the batch and frames dimensions
+        # sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width]
+        sample = sample.flatten(0, 1)
+        # Repeat the embeddings num_video_frames times
+        # emb: [batch, channels] -> [batch * frames, channels]
+        emb = emb.repeat_interleave(num_frames, dim=0)
+        # encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels]
+        encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0)
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device)
+
+        # Branch: 4 down_blocks, each with 3 skip connections. Here we ignore the first skip branch, whose computations only has up_blocks but without down_blocks.
+        if cache_branch is not None:
+            each_module_num = len(self.down_blocks[0].resnets) + 1
+            down_cache_block_idx = cache_branch // each_module_num
+            down_cache_module_idx = cache_branch % each_module_num
+
+            up_cache_block_idx = len(self.up_blocks) - 1 - down_cache_block_idx
+            up_cache_module_idx = 1 - down_cache_module_idx
+            if down_cache_module_idx == each_module_num - 1:
+                up_cache_block_idx -= 1
+                up_cache_module_idx = 2
+            
+        if cache_features is not None:
+            # 3. down
+            down_block_res_samples = (sample,)
+            for block_id, downsample_block in enumerate(self.down_blocks):
+                if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                        exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None
+                    )
+                else:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        image_only_indicator=image_only_indicator,
+                        exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None
+                    )
+                    
+                down_block_res_samples += res_samples
+                if down_cache_block_idx == block_id:
+                    break
+
+            # 4. no mid
+            sample = cache_features
+
+            # 5. up
+            for i, upsample_block in enumerate(self.up_blocks):
+                if i < up_cache_block_idx:
+                    continue
+                
+                if i == up_cache_block_idx:
+                    trunc_res_samples_len = len(upsample_block.resnets) - up_cache_module_idx
+                else:
+                    trunc_res_samples_len = len(upsample_block.resnets)
+                
+                res_samples = down_block_res_samples[-trunc_res_samples_len :]
+                down_block_res_samples = down_block_res_samples[: -trunc_res_samples_len]
+
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    sample, _ = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                        enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None
+                    )
+                else:
+                    sample, _ = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        image_only_indicator=image_only_indicator,
+                        enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None
+                    )
+        else:
+            # 3. down
+            down_block_res_samples = (sample,)
+            for downsample_block in self.down_blocks:
+                if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                    )
+                else:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        image_only_indicator=image_only_indicator,
+                    )
+
+                down_block_res_samples += res_samples
+
+            # 4. mid
+            sample = self.mid_block(
+                hidden_states=sample,
+                temb=emb,
+                encoder_hidden_states=encoder_hidden_states,
+                image_only_indicator=image_only_indicator,
+            )
+
+            # 5. up
+            for i, upsample_block in enumerate(self.up_blocks):
+                res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+                down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    sample, current_record_f = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                    )
+                else:
+                    sample, current_record_f = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        image_only_indicator=image_only_indicator,
+                    )
+
+                if cache_branch is not None and i == up_cache_block_idx:
+                    cache_features = current_record_f[up_cache_module_idx]
+
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        # 7. Reshape back to original shape
+        sample = sample.reshape(batch_size, num_frames, *sample.shape[1:])
+
+        if not return_dict:
+            return (sample, cache_features)
+
+        return UNetSpatioTemporalConditionOutput(sample=sample)

ixformer_sdk/contrib/comfy/unet_model_wrapper.py

@@ -0,0 +1,496 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import ixformer.functions as ixf_F
+
+
+def time_embed(t_emb, weight1, bias1, weight2, bias2):
+    # unet time_emd
+    # linear + silu + linear
+    emb = ixf_F.act_bias_mm(
+        t_emb, weight1, act_type="silu", bias=bias1, scale=1, trans_format="TN"
+    )
+    emb = ixf_F.act_bias_mm(
+        emb, weight2, act_type="none", bias=bias2, scale=1, trans_format="TN"
+    )
+    return emb
+
+
+def ixf_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05):
+    return ixf_F.layernorm(input, weight, bias, normalized_shape)
+
+
+def ixf_pt_scaled_dot_product_attention(
+    query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
+):
+    if (
+        not query.is_contiguous()
+        and query.transpose(1, 2).is_contiguous()
+        and key.transpose(1, 2).is_contiguous()
+        and value.transpose(1, 2).is_contiguous()
+        and attn_mask is None
+    ):
+
+        batch_size, head_num, seq_len_q, head_dim = query.shape
+        _, _, seq_len_k, _ = key.shape
+
+        query = query.transpose(1, 2).view(batch_size * seq_len_q, head_num, head_dim)
+        key = key.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim)
+        value = value.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim)
+
+        cu_seqlens_q = torch.arange(
+            0,
+            seq_len_q * (batch_size + 1),
+            seq_len_q,
+            dtype=torch.int32,
+            device=query.device,
+        )
+        if seq_len_q == seq_len_k:
+            cu_seqlens_k = cu_seqlens_q
+        else:
+            cu_seqlens_k = torch.arange(
+                0,
+                seq_len_k * (batch_size + 1),
+                seq_len_k,
+                dtype=torch.int32,
+                device=query.device,
+            )
+
+        res = ixf_F.flash_attn_varlen_func(
+            query,
+            key,
+            value,
+            cu_seqlens_q.int(),
+            cu_seqlens_k.int(),
+            seq_len_q,
+            seq_len_k,
+        )
+        res = res.view(batch_size, seq_len_q, head_num, head_dim).transpose(1, 2)
+        return res
+
+    if not query.is_contiguous():
+        query = query.contiguous()
+    if not key.is_contiguous():
+        key = key.contiguous()
+    if not value.is_contiguous():
+        value = value.contiguous()
+    return ixf_F.scaled_dot_product_attention(
+        query, key, value, attn_mask=attn_mask, is_causal=is_causal
+    )
+
+
+class UnetIxformerFunction:
+    def __init__(self) -> None:
+        self.ixf_linear = ixf_F.linear
+        self.pt_linear = F.linear
+        self.pt_layer_norm = F.layer_norm
+        self.pt_scaled_dot_product_attention = F.scaled_dot_product_attention
+
+    def __enter__(self):
+        F.linear = self.ixf_linear
+        F.layer_norm = ixf_layer_norm
+        F.scaled_dot_product_attention = ixf_pt_scaled_dot_product_attention
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        F.linear = self.pt_linear
+        F.layer_norm = self.pt_layer_norm
+        F.scaled_dot_product_attention = self.pt_scaled_dot_product_attention
+        if exc_tb is not None:
+            print(f"{exc_type} {exc_val}")
+            return False
+        return True
+
+
+def ForwardWrapper(fun):
+    def wrap(*args, **kwargs):
+        with UnetIxformerFunction() as w:
+            return fun(*args, **kwargs)
+
+    return wrap
+
+
+class IxformerComfyWrapper(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.is_ixf_wrapper = True
+
+
+class Conv2dNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module):
+        super().__init__()
+        module.weight.data = module.weight.permute(0, 2, 3, 1).contiguous()
+        module.bias.data = module.bias.float()
+        self.weight = module.weight.data
+        self.bias = module.bias.data
+        self.stride = module.stride
+        self.padding = module.padding
+        self.dilation = module.dilation
+        self.groups = module.groups
+
+    def forward(self, x):
+        h2 = ixf_F.conv2d(
+            x,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
+        return h2
+
+
+class ResBlockNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module) -> None:
+        super().__init__()
+        assert not module.updown
+        assert not module.use_scale_shift_norm
+        assert not module.skip_t_emb
+        assert not module.exchange_temb_dims
+
+        if isinstance(module.skip_connection, nn.Identity):
+            self.skip_connection = module.skip_connection
+        elif get_class_name(module.skip_connection) == "Conv2d":
+            self.skip_connection = Conv2dNhwcWrapper(module.skip_connection)
+        else:
+            raise NotImplementedError(
+                f"ResBlockNhwcWrapper support Conv2d or nn.Identity, but got {module.skip_connection}"
+            )
+
+        self.in_layers = module.in_layers
+        self.out_layers = module.out_layers
+        self.emb_layers = module.emb_layers
+        self.in_layers_conv = Conv2dNhwcWrapper(module.in_layers[2])
+        self.out_layers_conv = Conv2dNhwcWrapper(module.out_layers[3])
+
+    def forward(self, x, emb):
+        # x: nhwc
+        x1 = x
+        # print(x1.shape)
+        # fused group_norm silu
+        h = ixf_F.group_norm(
+            x1,  # nchw->nhwc
+            self.in_layers[0].num_groups,
+            self.in_layers[0].weight,
+            self.in_layers[0].bias,
+            format=False,
+            act_type=1,
+        )
+        h = self.in_layers_conv(h)
+
+        emb_out = self.emb_layers(emb)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+
+        h = h + emb_out.permute(0, 2, 3, 1)
+        # print(h.shape)
+        h = ixf_F.group_norm(
+            h,
+            self.out_layers[0].num_groups,
+            self.out_layers[0].weight,
+            self.out_layers[0].bias,
+            format=False,
+            act_type=1,
+        )
+
+        h = self.out_layers[2](h)
+        h = self.out_layers_conv(h)
+        # TODO: support other skip_connection
+        return self.skip_connection(x) + h
+
+
+class DownsampleNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module) -> None:
+        # TODO: support avg_pool_nd
+        super().__init__()
+        assert module.use_conv
+        self.channels = module.channels
+        self.op = Conv2dNhwcWrapper(module.op)
+
+    def forward(self, x):
+        assert x.shape[-1] == self.channels
+        return self.op(x)
+
+
+def ffn_forward(self, x):
+    if get_class_name(self.net[0]) == "GEGLU":
+        net = self.net[1:]
+        geglu_net = self.net[0]
+        x = geglu_net.proj(x)
+        x = ixf_F.gelu_and_mul(x)
+        return net(x)
+    else:
+        return self.net(x)
+
+
+# ComfyUI/comfy/ldm/modules/attention.py `class BasicTransformerBlock(nn.Module)`
+def transformer_block_forward(self, x, context=None, transformer_options={}):
+    extra_options = {}
+    block = transformer_options.get("block", None)
+    block_index = transformer_options.get("block_index", 0)
+    transformer_patches = {}
+    transformer_patches_replace = {}
+
+    for k in transformer_options:
+        if k == "patches":
+            transformer_patches = transformer_options[k]
+        elif k == "patches_replace":
+            transformer_patches_replace = transformer_options[k]
+        else:
+            extra_options[k] = transformer_options[k]
+
+    extra_options["n_heads"] = self.n_heads
+    extra_options["dim_head"] = self.d_head
+
+    if self.ff_in:
+        x_skip = x
+        x = self.ff_in(self.norm_in(x))
+        if self.is_res:
+            x += x_skip
+
+    n = self.norm1(x)
+    if self.disable_self_attn:
+        context_attn1 = context
+    else:
+        context_attn1 = None
+    value_attn1 = None
+
+    if "attn1_patch" in transformer_patches:
+        patch = transformer_patches["attn1_patch"]
+        if context_attn1 is None:
+            context_attn1 = n
+        value_attn1 = context_attn1
+        for p in patch:
+            n, context_attn1, value_attn1 = p(
+                n, context_attn1, value_attn1, extra_options
+            )
+
+    if block is not None:
+        transformer_block = (block[0], block[1], block_index)
+    else:
+        transformer_block = None
+
+    attn1_replace_patch = transformer_patches_replace.get("attn1", {})
+    block_attn1 = transformer_block
+    if block_attn1 not in attn1_replace_patch:
+        block_attn1 = block
+
+    if block_attn1 in attn1_replace_patch:
+        if context_attn1 is None:
+            context_attn1 = n
+            value_attn1 = n
+        n = self.attn1.to_q(n)
+        context_attn1 = self.attn1.to_k(context_attn1)
+        value_attn1 = self.attn1.to_v(value_attn1)
+        n = attn1_replace_patch[block_attn1](
+            n, context_attn1, value_attn1, extra_options
+        )
+        n = self.attn1.to_out(n)
+    else:
+        n = self.attn1(n, context=context_attn1, value=value_attn1)
+
+    if "attn1_output_patch" in transformer_patches:
+        patch = transformer_patches["attn1_output_patch"]
+        for p in patch:
+            n = p(n, extra_options)
+
+    x += n
+    if "middle_patch" in transformer_patches:
+        patch = transformer_patches["middle_patch"]
+        for p in patch:
+            x = p(x, extra_options)
+
+    if self.attn2 is not None:
+        n = self.norm2(x)
+        if self.switch_temporal_ca_to_sa:
+            context_attn2 = n
+        else:
+            context_attn2 = context
+        value_attn2 = None
+        if "attn2_patch" in transformer_patches:
+            patch = transformer_patches["attn2_patch"]
+            value_attn2 = context_attn2
+            for p in patch:
+                n, context_attn2, value_attn2 = p(
+                    n, context_attn2, value_attn2, extra_options
+                )
+
+        attn2_replace_patch = transformer_patches_replace.get("attn2", {})
+        block_attn2 = transformer_block
+        if block_attn2 not in attn2_replace_patch:
+            block_attn2 = block
+
+        if block_attn2 in attn2_replace_patch:
+            if value_attn2 is None:
+                value_attn2 = context_attn2
+            n = self.attn2.to_q(n)
+            context_attn2 = self.attn2.to_k(context_attn2)
+            value_attn2 = self.attn2.to_v(value_attn2)
+            n = attn2_replace_patch[block_attn2](
+                n, context_attn2, value_attn2, extra_options
+            )
+            n = self.attn2.to_out(n)
+        else:
+            n = self.attn2(n, context=context_attn2, value=value_attn2)
+
+    if "attn2_output_patch" in transformer_patches:
+        patch = transformer_patches["attn2_output_patch"]
+        for p in patch:
+            n = p(n, extra_options)
+
+    # x += n
+    # if self.is_res:
+    #     x_skip = x
+    # x = self.ff(self.norm3(x))
+
+    x, x_skip = ixf_F.residual_layer_norm(
+        n,
+        self.norm3.normalized_shape,
+        self.norm3.weight,
+        self.norm3.bias,
+        x,
+        eps=self.norm3.eps,
+        is_post_ln=False,
+    )
+    x = ffn_forward(self.ff, x)
+
+    # x = ffn_forward(self.ff, self.norm3(x))
+    if self.is_res:
+        x += x_skip
+
+    return x
+
+
+class SpatialTransformerNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module):
+        super().__init__()
+        self.use_linear = module.use_linear
+        self.transformer_blocks = module.transformer_blocks
+        self.norm = module.norm
+        if not self.use_linear:
+            self.proj_in = Conv2dNhwcWrapper(module.proj_in)
+            self.proj_out = Conv2dNhwcWrapper(module.proj_out)
+        else:
+            self.proj_in = module.proj_in
+            self.proj_out = module.proj_out
+
+    @ForwardWrapper
+    def forward(self, x, context=None, transformer_options={}):
+        # note: if no context is given, cross-attention defaults to self-attention
+        if not isinstance(context, list):
+            context = [context] * len(self.transformer_blocks)
+
+        b, h, w, c = x.shape
+        x_in = x
+
+        # group_norm
+        x = ixf_F.group_norm(
+            x,
+            self.norm.num_groups,
+            self.norm.weight,
+            self.norm.bias,
+            format=False,
+        )
+        # conv2d
+        if not self.use_linear:
+            x = self.proj_in(x)
+        # n,(hw),c
+        x = x.view(x.shape[0], -1, x.shape[-1])
+        if self.use_linear:
+            x = self.proj_in(x)
+
+        for i, block in enumerate(self.transformer_blocks):
+            transformer_options["block_index"] = i
+            # x = block(x, context=context[i], transformer_options=transformer_options)
+            x = transformer_block_forward(
+                block, x, context=context[i], transformer_options=transformer_options
+            )
+
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = x.view(b, h, w, c)
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in
+
+
+class UpsampleNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module) -> None:
+        # TODO: support mhwc interpolate
+        super().__init__()
+        self.dims = module.dims
+        self.use_conv = module.use_conv
+        self.channels = module.channels
+        if self.use_conv:
+            self.conv = Conv2dNhwcWrapper(module.conv)
+
+    def forward(self, x, output_shape=None):
+        # print("================== Upsample is running ==================")
+        assert x.shape[-1] == self.channels
+        assert len(x.shape) == 4
+
+        # nhwc -> nchw
+        if output_shape is not None:
+            assert len(output_shape) == 4
+            output_shape = [
+                output_shape[0],
+                output_shape[3],
+                output_shape[1],
+                output_shape[2],
+            ]
+        x = x.permute(0, 3, 1, 2).contiguous()
+        if self.dims == 3:
+            shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2]
+            if output_shape is not None:
+                shape[1] = output_shape[3]
+                shape[2] = output_shape[4]
+        else:
+            shape = [x.shape[2] * 2, x.shape[3] * 2]
+            if output_shape is not None:
+                shape[0] = output_shape[2]
+                shape[1] = output_shape[3]
+        # TODO: interpolate 支持 nhwc， 去掉前后转置
+        x = F.interpolate(x, size=shape, mode="nearest")
+        # nchw -> nhwc
+        x = x.permute(0, 2, 3, 1).contiguous()
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+unet_wrappers = {
+    "Conv2d": Conv2dNhwcWrapper,
+    "ResBlock": ResBlockNhwcWrapper,
+    "Downsample": DownsampleNhwcWrapper,
+    "SpatialTransformer": SpatialTransformerNhwcWrapper,
+    "Upsample": UpsampleNhwcWrapper,
+}
+
+
+def get_class_name(module):
+    return module.__class__.__name__
+
+
+def module_wrapper(module):
+    # 将原始的 module 封装为 nhwc 模式
+    module_name = get_class_name(module)
+    assert (
+        module_name == "TimestepEmbedSequential"
+    ), f"ixformer unet_model_wrapper only support 'TimestepEmbedSequential' now, but got {module_name}"
+
+    num_sequential = len(module)
+    for idx_seq in range(num_sequential):
+        sub_module = module[idx_seq]
+        sub_module_name = get_class_name(sub_module)
+        # 判断模块是否已经封装
+        if not getattr(sub_module, "is_ixf_wrapper", False):
+            if sub_module_name in unet_wrappers:
+                module[idx_seq].forward = unet_wrappers[sub_module_name](
+                    sub_module
+                ).forward
+                module[idx_seq].is_ixf_wrapper = True
+            else:
+                raise NotImplementedError(f"{sub_module_name} not support")
+    return module

ixformer_sdk/contrib/flashinfer/__init__.py

@@ -0,0 +1,17 @@
+from .decode import BatchDecodeWithPagedKVCacheWrapper
+from .prefill import (
+    BatchPrefillWithPagedKVCacheWrapper,
+    BatchPrefillWithRaggedKVCacheWrapper,
+)
+
+
+def bmm_fp8():
+    pass
+
+
+def SegmentGEMMWrapper():
+    pass
+
+
+def bmm_fp8():
+    pass

ixformer_sdk/contrib/flashinfer/activation.py

@@ -0,0 +1,29 @@
+import ixformer.inference.functions as ops
+import torch
+
+
+def gelu_and_mul():
+    pass
+
+
+def gelu_tanh_and_mul():
+    pass
+
+
+def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
+    r"""Fused SiLU and Mul operation.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (..., 2 * hidden_size).
+
+    out: Optional[torch.Tensor]
+        The the output tensor, if specified, the kernel will update this tensor inplace.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Output tensor, shape (..., hidden_size).
+    """
+    return ops.silu_and_mul(input=input, output=out)

ixformer_sdk/contrib/flashinfer/cascade.py

@@ -0,0 +1,2 @@
+def merge_state():
+    pass

ixformer_sdk/contrib/flashinfer/decode.py

@@ -0,0 +1,101 @@
+import math
+from typing import Optional, Tuple, Union
+
+import ixformer.inference.functions as ops
+import torch
+
+
+def _grouped_size_compiled_for_decode_kernels(
+    num_qo_heads: int, num_kv_heads: int
+) -> bool:
+    return (num_qo_heads // num_kv_heads) in [1, 2, 4, 8]
+
+
+class BatchDecodeWithPagedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+        use_cuda_graph: bool = False,
+        use_tensor_cores: bool = False,
+    ) -> None:
+        pass
+
+    def plan(
+        self,
+        indptr: torch.Tensor,
+        indices: torch.Tensor,
+        last_page_len: torch.Tensor,
+        num_qo_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        page_size: int,
+        # pos_encoding_mode: str = "NONE",
+        # window_left: int = -1,
+        # logits_soft_cap: Optional[float] = None,
+        data_type: Union[str, torch.dtype] = "float16",
+        q_data_type: Optional[Union[str, torch.dtype]] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+        max_seqlen_q: int = None,
+        max_seqlen_k: int = None,
+    ) -> None:
+        self.indptr = indptr
+        self.indices = indices
+        self.last_page_len = last_page_len
+        self.num_qo_heads = num_qo_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+
+        assert page_size == 1
+
+        self.cu_seqlens_q = torch.ones_like(indptr)
+        self.cu_seqlens_q[0] = 0
+        self.cu_seqlens_q = torch.cumsum(self.cu_seqlens_q, dim=0).int()
+
+        self.cu_seqlens_k = indptr
+        if sm_scale is None:
+            sm_scale = 1.0 / math.sqrt(head_dim)
+
+            self.sm_scale = sm_scale
+            self.max_seqlen_q = max_seqlen_q
+            self.max_seqlen_k = max_seqlen_k
+
+    begin_forward = plan
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+        pos_encoding_mode: str = "NONE",
+        q_scale: Optional[float] = None,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        window_left: int = -1,
+        logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        rope_scale: Optional[float] = None,
+        rope_theta: Optional[float] = None,
+    ) -> torch.Tensor:
+        k_cache, v_cache = paged_kv_cache
+
+        out = torch.empty_like(q)
+
+        ops.paged_attention_flashinfer(
+            output=out,
+            query=q,
+            paged_kv_data=(k_cache.unsqueeze(1), v_cache.unsqueeze(1)),
+            paged_kv_indptr=self.indptr,
+            paged_kv_indices=self.indices,
+            paged_kv_last_page_len=self.last_page_len,
+            scale=self.sm_scale,
+            max_seq_len=self.max_seqlen_k,
+            kv_cache_format="NHD",
+        )
+
+        return out
+
+    def end_forward(self) -> None:
+        r"""Warning: this function is deprecated and has no effect."""
+        pass

ixformer_sdk/contrib/flashinfer/norm.py

@@ -0,0 +1,61 @@
+import ixformer.inference.functions as ops
+import torch
+
+
+def fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+):
+    r"""Fused add root mean square normalization.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    residual: torch.Tensor
+        Residual tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+    """
+    return ops.residual_rms_norm(
+        input=input,
+        residual=residual,
+        weight=weight,
+        eps=eps,
+    )
+
+
+def gemma_fused_add_rmsnorm():
+    pass
+
+
+def gemma_rmsnorm():
+    pass
+
+
+def rmsnorm(
+    input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    r"""Root mean square normalization.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Normalized tensor, shape (batch_size, hidden_size).
+    """
+
+    return ops.rms_norm(
+        input=input,
+        weight=weight,
+        eps=eps,
+    )

ixformer_sdk/contrib/flashinfer/prefill.py

@@ -0,0 +1,113 @@
+import math
+from typing import Optional, Tuple, Union
+
+import ixformer._C as ops
+import torch
+
+
+class BatchPrefillWithRaggedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+    ):
+        pass
+
+    def plan(
+        self,
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        num_qo_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        # custom_mask: Optional[torch.Tensor] = None,
+        # packed_custom_mask: Optional[torch.Tensor] = None,
+        causal: bool = True,
+        # pos_encoding_mode: str = "NONE",
+        # allow_fp16_qk_reduction: bool = False,
+        # window_left: int = -1,
+        # logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+        # q_data_type: str = "float16",
+    ) -> None:
+        batch_size = len(qo_indptr) - 1
+        if len(kv_indptr) != batch_size + 1:
+            raise ValueError(
+                "The kv_indptr length should be equal to qk_indptr length."
+            )
+        self._causal = causal
+        self._sm_scale = sm_scale
+        if sm_scale is None:
+            sm_scale = 1.0 / math.sqrt(head_dim)
+
+        self.cu_seqlens_q = qo_indptr
+        self.cu_seqlens_k = kv_indptr
+        self.num_qo_heads = num_qo_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.max_seqlen_q = max_seqlen_q
+        self.max_seqlen_k = max_seqlen_k
+
+    begin_forward = plan
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        causal: bool = True,
+        # pos_encoding_mode: str = "NONE",
+        # allow_fp16_qk_reduction: bool = False,
+        # window_left: int = -1,
+        logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+    ) -> torch.Tensor:
+        r"""Warning: This function is deprecated, please use :meth:`run` instead."""
+
+        q = q.view(-1, self.num_qo_heads, self.head_dim)
+        k = k.view(-1, self.num_kv_heads, self.head_dim)
+        v = v.view(-1, self.num_kv_heads, self.head_dim)
+
+        out = torch.empty_like(q)
+
+        assert causal
+        assert (
+            logits_soft_cap is None or logits_soft_cap == 0
+        ), f"logits_soft_cap not supported, but got logits_soft_cap={logits_soft_cap}"
+
+        ops.infer.ixinfer_flash_attn_unpad(
+            q,
+            k,
+            v,
+            out,
+            self.cu_seqlens_q,
+            self.cu_seqlens_k,
+            self.max_seqlen_q,
+            self.max_seqlen_k,
+            causal,
+            False,  # need_lse =False
+            sm_scale,
+            False,
+            None,
+        )
+        return out
+
+    def end_forward(self) -> None:
+        r"""Warning: this function is deprecated and has no effect."""
+        pass
+
+
+class BatchPrefillWithPagedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+        use_cuda_graph: bool = False,
+    ) -> None:
+        pass

ixformer_sdk/contrib/flashinfer/sampling.py

@@ -0,0 +1,14 @@
+def min_p_sampling_from_probs():
+    pass
+
+
+def top_k_renorm_prob():
+    pass
+
+
+def top_k_top_p_sampling_from_probs():
+    pass
+
+
+def top_p_renorm_prob():
+    pass

ixformer_sdk/contrib/tgi/__init__.py

@@ -0,0 +1 @@
+from .fused_moe import fused_moe

ixformer_sdk/contrib/tgi/fused_moe.py

@@ -0,0 +1,430 @@
+import functools
+import json
+import os
+from typing import Any, Dict, Optional, Tuple
+from loguru import logger
+import torch
+import ixformer.inference.functions as ops
+
+CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "65536"))
+
+def fused_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], (
+        "Number of tokens mismatch")
+
+    M, _ = hidden_states.shape
+
+    topk_weights = torch.empty(M,
+                               topk,
+                               dtype=torch.float32,
+                               device=hidden_states.device)
+    topk_ids = torch.empty(M,
+                           topk,
+                           dtype=torch.int32,
+                           device=hidden_states.device)
+    token_expert_indicies = torch.empty(M,
+                                        topk,
+                                        dtype=torch.int32,
+                                        device=hidden_states.device)
+    ops.vllm_moe_topk_softmax(
+        topk_weights,
+        topk_ids,
+        token_expert_indicies,
+        gating_output.float(),  # TODO(woosuk): Optimize this.
+    )
+    del token_expert_indicies  # Not used. Will be used in the future.
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+    return topk_weights, topk_ids
+
+# This is used by the Deepseek-V2 model
+def grouped_topk(hidden_states: torch.Tensor,
+                 gating_output: torch.Tensor,
+                 topk: int,
+                 renormalize: bool,
+                 num_expert_group: int = 0,
+                 topk_group: int = 0):
+
+    assert hidden_states.shape[0] == gating_output.shape[0], (
+        "Number of tokens mismatch")
+
+    scores = torch.softmax(gating_output, dim=-1)
+    num_token = scores.shape[0]
+    group_scores = scores.view(num_token, num_expert_group,
+                               -1).max(dim=-1).values  # [n, n_group]
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
+                           sorted=False)[1]  # [n, top_k_group]
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = group_mask.unsqueeze(-1).expand(
+        num_token, num_expert_group,
+        scores.shape[-1] // num_expert_group).reshape(num_token, -1)  # [n, e]
+    tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0)  # [n, e]
+    topk_weights, topk_ids = torch.topk(tmp_scores,
+                                        k=topk,
+                                        dim=-1,
+                                        sorted=False)
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+    return topk_weights, topk_ids
+
+def get_config_file_name(E: int, N: int, dtype: Optional[str]) -> str:
+    device_name = torch.cuda.get_device_name().replace(" ", "_")
+    dtype_selector = "" if not dtype else f",dtype={dtype}"
+    return f"E={E},N={N},device_name={device_name}{dtype_selector}.json"
+
+@functools.lru_cache
+def get_moe_configs(E: int, N: int,
+                    dtype: Optional[str]) -> Optional[Dict[int, Any]]:
+    """
+    Return optimized configurations for the fused MoE kernel.
+
+    The return value will be a dictionary that maps an irregular grid of
+    batch sizes to configurations of the fused_moe kernel. To evaluate the
+    kernel on a given batch size bs, the closest batch size in the grid should
+    be picked and the associated configuration chosen to invoke the kernel.
+    """
+
+    # First look up if an optimized configuration is available in the configs
+    # directory
+    json_file_name = get_config_file_name(E, N, dtype)
+
+    config_file_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
+    if os.path.exists(config_file_path):
+        with open(config_file_path) as f:
+            logger.info("Using configuration from %s for MoE layer.",
+                        config_file_path)
+            # If a configuration has been found, return it
+            return {int(key): val for key, val in json.load(f).items()}
+
+    # If no optimized configuration is available, we will use the default
+    # configuration
+    return None
+
+def get_default_config(
+    M: int,
+    E: int,
+    N: int,
+    K: int,
+    topk: int,
+    dtype: Optional[str],
+) -> Dict[str, int]:
+    config = {
+        'BLOCK_SIZE_M': 64,
+        'BLOCK_SIZE_N': 64,
+        'BLOCK_SIZE_K': 32,
+        'GROUP_SIZE_M': 8
+    }
+    if M <= E:
+        config = {
+            'BLOCK_SIZE_M': 16,
+            'BLOCK_SIZE_N': 32,
+            'BLOCK_SIZE_K': 64,
+            'GROUP_SIZE_M': 1
+        }
+    numel = M * topk
+    if numel <= 64:
+        config['BLOCK_SIZE_M'] = 32
+    elif numel <= 1024:
+        config['BLOCK_SIZE_M'] = 64
+    else:
+        config['BLOCK_SIZE_M'] = 256
+    return config
+
+def try_get_optimal_moe_config(
+    w1_shape: Tuple[int, ...],
+    w2_shape: Tuple[int, ...],
+    top_k: int,
+    dtype: Optional[str],
+    M: int,
+    override_config: Optional[Dict[str, Any]] = None,
+):
+    if override_config:
+        config = override_config
+    else:
+        # First try to load optimal config from the file
+        E, _, N = w2_shape
+        configs = get_moe_configs(E, N, dtype)
+
+        if configs:
+            # If an optimal configuration map has been found, look up the
+            # optimal config
+            config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
+        else:
+            # Else use the default config
+            config = get_default_config(M, E, N, w1_shape[2], top_k, dtype)
+    return config
+
+def moe_align_block_size(
+        topk_ids: torch.Tensor, block_size: int,
+        num_experts: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Aligns the token distribution across experts to be compatible with block
+    size for matrix multiplication.
+
+    Parameters:
+    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
+        top-k expert indices for each token.
+    - block_size: The block size used in block matrix multiplication.
+    - num_experts: The total number of experts.
+
+    Returns:
+    - sorted_token_ids: A tensor containing the sorted token indices according
+        to their allocated expert.
+    - expert_ids: A tensor indicating the assigned expert index for each block.
+    - num_tokens_post_padded: The total number of tokens after padding,
+        ensuring divisibility by block_size.
+
+    This function pads the number of tokens that each expert needs to process
+    so that it is divisible by block_size.
+    Padding ensures that during block matrix multiplication, the dimensions
+    align correctly.
+
+    Example:
+    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
+    block_size = 4, and num_experts = 4:
+    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
+        with each expert needing to process 3 tokens.
+    - As block_size is 4, we pad 1 token for each expert.
+    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
+    - Then append padding tokens [12, 12, 12, 12] for each block.
+    - After sorting by expert index, we obtain token_ids
+        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
+        Tokens 12 are non-existent (padding) and are ignored in
+        the subsequent matrix multiplication.
+    - The padding ensures that the total number of tokens is now divisible
+        by block_size for proper block matrix operations.
+    """
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    sorted_ids = torch.empty((max_num_tokens_padded, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    sorted_ids.fill_(topk_ids.numel())
+     # max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
+    max_num_m_blocks = topk_ids.numel() + num_experts
+    expert_ids = torch.empty((max_num_m_blocks, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    num_tokens_post_pad = torch.empty((1),
+                                      dtype=torch.int32,
+                                      device=topk_ids.device)
+    ops.vllm_moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
+                             expert_ids, num_tokens_post_pad)
+    return sorted_ids, expert_ids, num_tokens_post_pad
+
+def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
+                            A_scale: Optional[torch.Tensor],
+                            B_scale: Optional[torch.Tensor],
+                            topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+                            sorted_token_ids: torch.Tensor,
+                            expert_ids: torch.Tensor,
+                            num_tokens_post_padded: torch.Tensor,
+                            mul_routed_weight: bool, top_k: int,
+                            config: Dict[str, Any], compute_type: torch.dtype,
+                            use_fp8: bool) -> None:
+    ops.vllm_invoke_fused_moe_kernel(A, B, C, topk_weights, topk_ids,
+                                     sorted_token_ids,expert_ids, num_tokens_post_padded,
+                                     mul_routed_weight, top_k, config['BLOCK_SIZE_M'])
+
+def fused_experts(hidden_states: torch.Tensor,
+                  w1: torch.Tensor,
+                  w2: torch.Tensor,
+                  topk_weights: torch.Tensor,
+                  topk_ids: torch.Tensor,
+                  inplace: bool = False,
+                  override_config: Optional[Dict[str, Any]] = None,
+                  use_fp8: bool = False,
+                  w1_scale: Optional[torch.Tensor] = None,
+                  w2_scale: Optional[torch.Tensor] = None,
+                  a1_scale: Optional[torch.Tensor] = None,
+                  a2_scale: Optional[torch.Tensor] = None):
+    # Check constraints.
+    assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.is_contiguous(), "Expert weights1 must be contiguous"
+    assert w2.is_contiguous(), "Expert weights2 must be contiguous"
+    assert hidden_states.dtype in [
+        torch.float32, torch.float16, torch.bfloat16
+    ]
+
+    num_tokens, _ = hidden_states.shape
+    E, N, _ = w1.shape
+    # We execute the fused_moe kernel in chunks to circumvent this issue:
+    # https://github.com/vllm-project/vllm/issues/5938
+    M = min(num_tokens, CHUNK_SIZE)
+
+    get_config_func = functools.partial(
+        try_get_optimal_moe_config,
+        w1.shape,
+        w2.shape,
+        topk_ids.shape[1],
+        "float8" if use_fp8 else None,
+        override_config=override_config,
+    )
+
+    config = get_config_func(M)
+
+    intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
+                                      device=hidden_states.device,
+                                      dtype=hidden_states.dtype)
+    intermediate_cache2 = torch.empty((M * topk_ids.shape[1], N // 2),
+                                      device=hidden_states.device,
+                                      dtype=hidden_states.dtype)
+    intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]),
+                                      device=hidden_states.device,
+                                      dtype=hidden_states.dtype)
+
+    compute_type = (torch.bfloat16
+                    if hidden_states.dtype == torch.bfloat16 else torch.float16)
+
+    if inplace:
+        out_hidden_states = hidden_states
+    else:
+        out_hidden_states = torch.empty_like(hidden_states)
+
+    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+        begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
+                                          min((chunk + 1) * CHUNK_SIZE,
+                                              num_tokens))
+        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+        tokens_in_chunk, _ = curr_hidden_states.shape
+
+        if tokens_in_chunk == 0:
+            break
+
+        if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
+            # Adjust the intermediate cache size and config for the last
+            # chunk. Note that in most cases we only have one chunk
+            # so the cache size and config are already set correctly and
+            # do not need to be adjusted.
+            intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
+            intermediate_cache2 = intermediate_cache2[:tokens_in_chunk]
+            intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
+            config = get_config_func(tokens_in_chunk)
+
+        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+        sorted_token_ids, expert_ids, num_tokens_post_padded = (
+            moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'], E))
+
+        invoke_fused_moe_kernel(curr_hidden_states,
+                                w1,
+                                intermediate_cache1,
+                                a1_scale,
+                                w1_scale,
+                                curr_topk_weights,
+                                curr_topk_ids,
+                                sorted_token_ids,
+                                expert_ids,
+                                num_tokens_post_padded,
+                                False,
+                                topk_ids.shape[1],
+                                config,
+                                compute_type=compute_type,
+                                use_fp8=use_fp8)
+
+        ops.silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
+
+        invoke_fused_moe_kernel(intermediate_cache2,
+                                w2,
+                                intermediate_cache3,
+                                a2_scale,
+                                w2_scale,
+                                curr_topk_weights,
+                                curr_topk_ids,
+                                sorted_token_ids,
+                                expert_ids,
+                                num_tokens_post_padded,
+                                True,
+                                1,
+                                config,
+                                compute_type=compute_type,
+                                use_fp8=use_fp8)
+
+        torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
+                  dim=1,
+                  out=out_hidden_states[begin_chunk_idx:end_chunk_idx])
+    return out_hidden_states
+
+def fused_moe(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    inplace: bool = False,
+    override_config: Optional[Dict[str, Any]] = None,
+    use_grouped_topk: bool = False,
+    num_expert_group: Optional[int] = None,
+    topk_group: Optional[int] = None,
+    use_fp8: bool = False,
+    w1_scale: Optional[torch.Tensor] = None,
+    w2_scale: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    This function computes a Mixture of Experts (MoE) layer using two sets of
+    weights, w1 and w2, and top-k gating mechanism.
+
+    Parameters:
+    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
+    - w1 (torch.Tensor): The first set of expert weights.
+    - w2 (torch.Tensor): The second set of expert weights.
+    - gating_output (torch.Tensor): The output of the gating operation
+        (before softmax).
+    - topk (int): The number of top-k experts to select.
+    - renormalize (bool): If True, renormalize the top-k weights to sum to 1.
+    - inplace (bool): If True, perform the operation in-place.
+        Defaults to False.
+    - override_config (Optional[Dict[str, Any]]): Optional override
+        for the kernel configuration.
+    - num_expert_group: Optional[int]: additional parameter for grouped_topk
+    - topk_group: Optional[int]: additional parameter for grouped_topk
+    - use_grouped_topk: If True, use grouped_topk instead of fused_topk
+        note: Deepseekv2 model uses grouped_topk
+    - use_fp8 (bool): If True, use fp8 arithmetic to compute the inner
+        products for w1 and w2. Defaults to False.
+    - w1_scale (Optional[torch.Tensor]): Optional scale to be used for
+        w1.
+    - w2_scale (Optional[torch.Tensor]): Optional scale to be used for
+        w2.
+
+    Returns:
+    - torch.Tensor: The output tensor after applying the MoE layer.
+    """
+    # Check constraints.
+    assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
+
+    if use_grouped_topk:
+        assert num_expert_group is not None and topk_group is not None
+        topk_weights, topk_ids = grouped_topk(hidden_states, gating_output,
+                                              topk, renormalize,
+                                              num_expert_group, topk_group)
+    else:
+        topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
+                                            renormalize)
+
+    return fused_experts(hidden_states,
+                         w1,
+                         w2,
+                         topk_weights,
+                         topk_ids,
+                         inplace=inplace,
+                         override_config=override_config,
+                         use_fp8=use_fp8,
+                         w1_scale=w1_scale,
+                         w2_scale=w2_scale,
+                         a1_scale=a1_scale,
+                         a2_scale=a2_scale)

ixformer_sdk/contrib/transformers/__init__.py

@@ -0,0 +1,2 @@
+from .models.bert.modeling_bert import BertForQuestionAnswering
+from .models.t5.modeling_t5 import T5ForConditionalGeneration

ixformer_sdk/contrib/transformers/models/bert/configuration_bert.py

@@ -0,0 +1,150 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     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.
+"""BERT model configuration"""
+
+from collections import OrderedDict
+from typing import Mapping
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.onnx import OnnxConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class BertConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`BertModel`] or a [`TFBertModel`]. It is used to
+    instantiate a BERT model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the BERT
+    [google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 30522):
+            Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`BertModel`] or [`TFBertModel`].
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+        hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"silu"` and `"gelu_new"` are supported.
+        hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (`int`, *optional*, defaults to 512):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        type_vocab_size (`int`, *optional*, defaults to 2):
+            The vocabulary size of the `token_type_ids` passed when calling [`BertModel`] or [`TFBertModel`].
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
+            Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
+            positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
+            [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
+            For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
+            with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
+        is_decoder (`bool`, *optional*, defaults to `False`):
+            Whether the model is used as a decoder or not. If `False`, the model is used as an encoder.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        classifier_dropout (`float`, *optional*):
+            The dropout ratio for the classification head.
+
+    Examples:
+
+    ```python
+    >>> from transformers import BertConfig, BertModel
+
+    >>> # Initializing a BERT google-bert/bert-base-uncased style configuration
+    >>> configuration = BertConfig()
+
+    >>> # Initializing a model (with random weights) from the google-bert/bert-base-uncased style configuration
+    >>> model = BertModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "bert"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=2,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        position_embedding_type="absolute",
+        use_cache=True,
+        classifier_dropout=None,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.position_embedding_type = position_embedding_type
+        self.use_cache = use_cache
+        self.classifier_dropout = classifier_dropout
+
+
+class BertOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        if self.task == "multiple-choice":
+            dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
+        else:
+            dynamic_axis = {0: "batch", 1: "sequence"}
+        return OrderedDict(
+            [
+                ("input_ids", dynamic_axis),
+                ("attention_mask", dynamic_axis),
+                ("token_type_ids", dynamic_axis),
+            ]
+        )

ixformer_sdk/contrib/transformers/models/t5/configuration_t5.py

@@ -0,0 +1,174 @@
+# coding=utf-8
+# Copyright 2020, The T5 Authors and HuggingFace Inc.
+#
+# 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.
+""" T5 model configuration"""
+from typing import Mapping
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.onnx import OnnxSeq2SeqConfigWithPast
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+T5_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "t5-small": "https://huggingface.co/t5-small/resolve/main/config.json",
+    "t5-base": "https://huggingface.co/t5-base/resolve/main/config.json",
+    "t5-large": "https://huggingface.co/t5-large/resolve/main/config.json",
+    "t5-3b": "https://huggingface.co/t5-3b/resolve/main/config.json",
+    "t5-11b": "https://huggingface.co/t5-11b/resolve/main/config.json",
+}
+
+
+class T5Config(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`T5Model`] or a [`TFT5Model`]. It is used to
+    instantiate a T5 model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the T5
+    [t5-small](https://huggingface.co/t5-small) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Arguments:
+        vocab_size (`int`, *optional*, defaults to 32128):
+            Vocabulary size of the T5 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`T5Model`] or [`TFT5Model`].
+        d_model (`int`, *optional*, defaults to 512):
+            Size of the encoder layers and the pooler layer.
+        d_kv (`int`, *optional*, defaults to 64):
+            Size of the key, query, value projections per attention head. The `inner_dim` of the projection layer will
+            be defined as `num_heads * d_kv`.
+        d_ff (`int`, *optional*, defaults to 2048):
+            Size of the intermediate feed forward layer in each `T5Block`.
+        num_layers (`int`, *optional*, defaults to 6):
+            Number of hidden layers in the Transformer encoder.
+        num_decoder_layers (`int`, *optional*):
+            Number of hidden layers in the Transformer decoder. Will use the same value as `num_layers` if not set.
+        num_heads (`int`, *optional*, defaults to 8):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        relative_attention_num_buckets (`int`, *optional*, defaults to 32):
+            The number of buckets to use for each attention layer.
+        relative_attention_max_distance (`int`, *optional*, defaults to 128):
+            The maximum distance of the longer sequences for the bucket separation.
+        dropout_rate (`float`, *optional*, defaults to 0.1):
+            The ratio for all dropout layers.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        initializer_factor (`float`, *optional*, defaults to 1):
+            A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+            testing).
+        feed_forward_proj (`string`, *optional*, defaults to `"relu"`):
+            Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. T5v1.1 uses the
+            `"gated-gelu"` feed forward projection. Original T5 uses `"relu"`.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+    """
+    model_type = "t5"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "hidden_size": "d_model",
+        "num_attention_heads": "num_heads",
+        "num_hidden_layers": "num_layers",
+    }
+
+    def __init__(
+        self,
+        vocab_size=32128,
+        d_model=512,
+        d_kv=64,
+        d_ff=2048,
+        num_layers=6,
+        num_decoder_layers=None,
+        num_heads=8,
+        relative_attention_num_buckets=32,
+        relative_attention_max_distance=128,
+        dropout_rate=0.1,
+        layer_norm_epsilon=1e-6,
+        initializer_factor=1.0,
+        feed_forward_proj="relu",
+        is_encoder_decoder=True,
+        use_cache=True,
+        pad_token_id=0,
+        eos_token_id=1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.d_kv = d_kv
+        self.d_ff = d_ff
+        self.num_layers = num_layers
+        self.num_decoder_layers = (
+            num_decoder_layers if num_decoder_layers is not None else self.num_layers
+        )  # default = symmetry
+        self.num_heads = num_heads
+        self.relative_attention_num_buckets = relative_attention_num_buckets
+        self.relative_attention_max_distance = relative_attention_max_distance
+        self.dropout_rate = dropout_rate
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_factor = initializer_factor
+        self.feed_forward_proj = feed_forward_proj
+        self.use_cache = use_cache
+
+        act_info = self.feed_forward_proj.split("-")
+        self.dense_act_fn = act_info[-1]
+        self.is_gated_act = act_info[0] == "gated"
+
+        if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2:
+            raise ValueError(
+                f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
+                "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
+                "'gated-gelu' or 'relu'"
+            )
+
+        # for backwards compatibility
+        if feed_forward_proj == "gated-gelu":
+            self.dense_act_fn = "gelu_new"
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            is_encoder_decoder=is_encoder_decoder,
+            **kwargs,
+        )
+
+
+class T5OnnxConfig(OnnxSeq2SeqConfigWithPast):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        common_inputs = {
+            "input_ids": {0: "batch", 1: "encoder_sequence"},
+            "attention_mask": {0: "batch", 1: "encoder_sequence"},
+        }
+        if self.use_past:
+            common_inputs["attention_mask"][1] = "past_encoder_sequence + sequence"
+            common_inputs["decoder_input_ids"] = {0: "batch"}
+            common_inputs["decoder_attention_mask"] = {
+                0: "batch",
+                1: "past_decoder_sequence + sequence",
+            }
+        else:
+            common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"}
+            common_inputs["decoder_attention_mask"] = {
+                0: "batch",
+                1: "decoder_sequence",
+            }
+
+        if self.use_past:
+            self.fill_with_past_key_values_(common_inputs, direction="inputs")
+
+        return common_inputs
+
+    @property
+    def default_onnx_opset(self) -> int:
+        return 13

ixformer_sdk/contrib/transformers/models/t5/ixformer_modeling_t5.py

@@ -0,0 +1,315 @@
+import torch
+from transformers.activations import NewGELUActivation
+
+import ixformer
+
+
+def self_attention_forward(
+    self,
+    hidden_states,
+    attention_mask=None,
+    position_bias=None,
+    layer_head_mask=None,
+    past_key_value=None,
+    use_cache=False,
+    output_attentions=False,
+):
+    assert output_attentions is False
+    assert layer_head_mask is None
+
+    normed_hidden_states = self.layer_norm(hidden_states)
+
+    if not hasattr(self, "qkv_weight"):
+        self.qkv_weight = torch.cat(
+            [
+                self.SelfAttention.q.weight,
+                self.SelfAttention.k.weight,
+                self.SelfAttention.v.weight,
+            ],
+            dim=0,
+        )
+        self.qkv_bias = None
+
+        del self.SelfAttention.q.weight
+        del self.SelfAttention.k.weight
+        del self.SelfAttention.v.weight
+
+    batch_size, seq_length = hidden_states.shape[:2]
+    real_seq_length = seq_length
+    if past_key_value is not None:
+        if len(past_key_value) != 2:
+            raise ValueError(
+                f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
+            )
+        real_seq_length += past_key_value[0].shape[2]
+    key_length = real_seq_length
+
+    def unshape(states):
+        """reshape"""
+        return (
+            states.transpose(1, 2)
+            .contiguous()
+            .view(batch_size, -1, self.SelfAttention.inner_dim)
+        )
+
+    qkv = ixformer.functions.linear(
+        normed_hidden_states, self.qkv_weight, self.qkv_bias
+    )
+
+    if past_key_value is not None:
+        pask_key, past_value = past_key_value
+        (
+            query_states,
+            key_states,
+            value_states,
+        ) = ixformer.functions.t5_split_qkv_update_kv_cache(
+            qkv,
+            pask_key,
+            past_value,
+            self.SelfAttention.n_heads,
+            self.SelfAttention.key_value_proj_dim,
+        )
+    else:
+        query_states, key_states, value_states = ixformer.functions.t5_split_qkv(
+            qkv, self.SelfAttention.n_heads, self.SelfAttention.key_value_proj_dim
+        )
+
+    if position_bias is None:
+        if not self.SelfAttention.has_relative_attention_bias:
+            position_bias = torch.zeros(
+                (1, self.SelfAttention.n_heads, real_seq_length, key_length),
+                device=query_states.device,
+                dtype=query_states.dtype,
+            )
+        else:
+            position_bias = self.SelfAttention.compute_bias(
+                real_seq_length, key_length, device=query_states.device
+            )
+
+        # if key and values are already calculated
+        # we want only the last query position bias
+        if past_key_value is not None:
+            position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
+
+        if attention_mask is not None:
+            # (batch_size, n_heads, seq_length, key_length)
+            position_bias = position_bias + attention_mask
+
+    if self.SelfAttention.pruned_heads:
+        mask = torch.ones(position_bias.shape[1])
+        mask[list(self.pruned_heads)] = 0
+        position_bias_masked = position_bias[:, mask.bool()]
+    else:
+        position_bias_masked = position_bias
+
+    attn_output = ixformer.functions.ixinfer_flash_attn_pad(
+        query_states.contiguous(),
+        key_states.contiguous(),
+        value_states.contiguous(),
+        mask=position_bias_masked.float().contiguous(),
+        atten_scale=1,
+    )
+    attn_output = unshape(attn_output)
+
+    attn_output = self.SelfAttention.o(attn_output)
+
+    present_key_value_state = (
+        (key_states, value_states)
+        if (self.SelfAttention.is_decoder and use_cache)
+        else None
+    )
+    outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)
+
+    if output_attentions:
+        outputs = outputs + (None,)
+    hidden_states = attn_output + hidden_states
+    outputs = (hidden_states,) + outputs[1:]
+
+    return outputs
+
+
+def cross_attention_forward(
+    self,
+    hidden_states,
+    key_value_states,
+    attention_mask=None,
+    position_bias=None,
+    layer_head_mask=None,
+    past_key_value=None,
+    use_cache=False,
+    query_length=None,
+    output_attentions=False,
+):
+
+    assert output_attentions is False
+    assert layer_head_mask is None
+
+    def unshape(states):
+        """reshape"""
+        return (
+            states.transpose(1, 2)
+            .contiguous()
+            .view(batch_size, -1, self.EncDecAttention.inner_dim)
+        )
+
+    normed_hidden_states = self.layer_norm(hidden_states)
+
+    # cross attn need key_value_states
+    assert key_value_states is not None
+    batch_size, seq_length = hidden_states.shape[:2]
+    real_seq_length = seq_length
+
+    if past_key_value is not None:
+        if len(past_key_value) != 2:
+            raise ValueError(
+                f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
+            )
+        real_seq_length += (
+            past_key_value[0].shape[2] if query_length is None else query_length
+        )
+
+    key_length = (
+        real_seq_length if key_value_states is None else key_value_states.shape[1]
+    )
+    head_num, head_dim = (
+        self.EncDecAttention.n_heads,
+        self.EncDecAttention.key_value_proj_dim,
+    )
+
+    query_states = (
+        self.EncDecAttention.q(normed_hidden_states)
+        .view(batch_size, seq_length, head_num, head_dim)
+        .transpose(1, 2)
+        .contiguous()
+    )
+
+    if past_key_value is not None:
+        if past_key_value[0].shape[2] != key_value_states.shape[1]:
+            # checking that the `sequence_length` of the `past_key_value` is the same as
+            # the provided `key_value_states` to support prefix tuning
+            # cross-attn
+            # (batch_size, n_heads, seq_length, dim_per_head)
+            key_states = (
+                self.EncDecAttention.k(key_value_states)
+                .view(batch_size, key_length, head_num, head_dim)
+                .transpose(1, 2)
+            )
+            value_states = (
+                self.EncDecAttention.v(key_value_states)
+                .view(batch_size, key_length, head_num, head_dim)
+                .transpose(1, 2)
+            )
+        else:
+            # cross-attn
+            key_states = past_key_value[0]
+            value_states = past_key_value[1]
+    else:
+        key_states = (
+            self.EncDecAttention.k(key_value_states)
+            .view(batch_size, key_length, head_num, head_dim)
+            .transpose(1, 2)
+        )
+        value_states = (
+            self.EncDecAttention.v(key_value_states)
+            .view(batch_size, key_length, head_num, head_dim)
+            .transpose(1, 2)
+        )
+
+    if not query_states.is_contiguous():
+        query_states = query_states.contiguous()
+
+    # TODO: fix this bug
+    if not value_states.is_contiguous():
+        new_value_states = query_states.new_empty(value_states.shape)
+        new_value_states.copy_(value_states)
+        value_states = new_value_states
+    if not key_states.is_contiguous():
+        numel = torch.numel(key_states)
+        new_key_states = query_states.new_empty([numel * 2])[:numel].view(
+            *list(key_states.shape)
+        )
+        new_key_states.copy_(key_states)
+        key_states = new_key_states
+
+    if position_bias is None:
+        if not self.EncDecAttention.has_relative_attention_bias:
+            position_bias = torch.zeros(
+                (1, self.EncDecAttention.n_heads, real_seq_length, key_length),
+                device=query_states.device,
+                dtype=query_states.dtype,
+            )
+        else:
+            position_bias = self.EncDecAttention.compute_bias(
+                real_seq_length, key_length, device=query_states.device
+            )
+
+        # if key and values are already calculated
+        # we want only the last query position bias
+        if past_key_value is not None:
+            position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
+
+        if attention_mask is not None:
+            # (batch_size, n_heads, seq_length, key_length)
+            position_bias = position_bias + attention_mask
+
+    if self.EncDecAttention.pruned_heads:
+        mask = torch.ones(position_bias.shape[1])
+        mask[list(self.pruned_heads)] = 0
+        position_bias_masked = position_bias[:, mask.bool()]
+    else:
+        position_bias_masked = position_bias
+
+    attn_output = ixformer.functions.ixinfer_flash_attn_pad(
+        query_states,
+        key_states.contiguous(),
+        value_states.contiguous(),
+        mask=position_bias_masked.float().contiguous(),
+        atten_scale=1,
+    )
+    attn_output = unshape(attn_output)
+
+    attn_output = self.EncDecAttention.o(attn_output)
+
+    present_key_value_state = (
+        (key_states, value_states)
+        if (self.EncDecAttention.is_decoder and use_cache)
+        else None
+    )
+    outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)
+
+    if output_attentions:
+        outputs = outputs + (None,)
+    hidden_states = attn_output + hidden_states
+    outputs = (hidden_states,) + outputs[1:]
+
+    return outputs
+
+
+def dense_gated_act_dense_forward(self, hidden_states):
+    if isinstance(self.act, NewGELUActivation):
+        if not hasattr(self, "wi"):
+            self.wi = torch.cat([self.wi_1.weight, self.wi_0.weight], dim=0)
+            del self.wi_1
+            del self.wi_0
+        hidden_states = ixformer.functions.linear(hidden_states, self.wi, None)
+        hidden_states = ixformer.functions.gelu_and_mul(hidden_states)
+        hidden_states = ixformer.functions.linear(hidden_states, self.wo.weight, None)
+    else:
+        hidden_gelu = self.act(self.wi_0(hidden_states))
+        hidden_linear = self.wi_1(hidden_states)
+        hidden_states = hidden_gelu * hidden_linear
+
+        hidden_states = self.dropout(hidden_states)
+
+        # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32.
+        # See https://github.com/huggingface/transformers/issues/20287
+        # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None``
+        if (
+            isinstance(self.wo.weight, torch.Tensor)
+            and hidden_states.dtype != self.wo.weight.dtype
+            and self.wo.weight.dtype != torch.int8
+        ):
+            hidden_states = hidden_states.to(self.wo.weight.dtype)
+
+        hidden_states = self.wo(hidden_states)
+    return hidden_states

ixformer_sdk/contrib/vllm/layers/__init__.py

@@ -0,0 +1,30 @@
+from .llama import forward_smoothquant
+from .mixtral import mixtral_decoder_layer_forward
+
+SUPPORT_REPLACE_METHOD = {
+    "llama": forward_smoothquant,
+}
+
+SUPPORT_REPLACE_LAYER = {
+    "llama": None,
+}
+
+
+def get_replace_forward(name: str):
+    try:
+        method = SUPPORT_REPLACE_METHOD[name]
+    except:
+        raise ValueError(
+            f"Only support replace names: {SUPPORT_REPLACE_METHOD.keys()}, but got {name}"
+        )
+    return method
+
+
+def get_replace_layer(name: str):
+    try:
+        layer = SUPPORT_REPLACE_LAYER[name]
+    except:
+        raise ValueError(
+            f"Only support replace names: {SUPPORT_REPLACE_LAYER.keys()}, but got {name}"
+        )
+    return layer

ixformer_sdk/contrib/vllm/layers/llama.py

@@ -0,0 +1,100 @@
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import torch
+from transformers import LlamaConfig
+
+from vllm.attention import AttentionMetadata
+from vllm.config import CacheConfig
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.models.llama import LlamaDecoderLayer as VllmLlamaDecoderLayer
+
+import vllm._custom_ops as ops
+# from ..overlap_comm import DecoderLayerOverlapComm, get_overlap_linear_method
+
+
+# This method is needed for support smoothquant no overlap forward
+def forward_smoothquant(
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    kv_caches: List[torch.Tensor],
+    attn_metadata: AttentionMetadata,
+    inputs_embeds: Optional[torch.Tensor] = None,
+    self = None, # will be set by partial
+) -> torch.Tensor:
+    dtype = self.dtype
+    
+    def forward_smoothquant_mlp(self,x,scales):
+        # gate_up_proj
+        # Int8 Matrix multiply.
+        bias = self.gate_up_proj.bias if not self.gate_up_proj.skip_bias_add else None
+        gate_up = ops.w8a8(x, self.gate_up_proj.weight, scales, self.gate_up_proj.weight_scales, dtype)
+        if bias:
+            gate_up += bias
+        
+        # act_fun
+        x, scales = ops.silu_and_mul_smoothquant(gate_up, self.down_proj.smooth_scales)
+        
+        # down_proj
+        output_parallel = ops.w8a8(x, self.down_proj.weight, scales, self.down_proj.weight_scales, dtype)
+        if self.down_proj.reduce_results and self.down_proj.tp_size > 1:
+            output = tensor_model_parallel_all_reduce(output_parallel)
+        else:
+            output = output_parallel
+
+        if not self.down_proj.skip_bias_add:
+            output = output + self.down_proj.bias if self.down_proj.bias is not None else output
+
+        return output
+    
+    def forward_smoothquant_attn(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        scales: torch.Tensor,
+    ) -> torch.Tensor:
+        # qkv proj
+        bias = self.qkv_proj.bias if not self.qkv_proj.skip_bias_add else None
+        
+        qkv = ops.w8a8(hidden_states, self.qkv_proj.weight, scales, self.qkv_proj.weight_scales, dtype)
+        if bias:
+            qkv += bias
+        
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
+        output, _ = self.o_proj(attn_output)  # TODO
+        return output
+    
+    if inputs_embeds is not None:
+        hidden_states = inputs_embeds
+    else:
+        hidden_states = self.get_input_embeddings(input_ids)
+    residual = None
+    for i in range(len(self.layers)):
+        layer = self.layers[i]
+        if residual is None:
+            residual = hidden_states
+            hidden_states, scales = ops.rms_norm_smoothquant(hidden_states,layer.input_layernorm.weight,layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
+        else:
+            hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.input_layernorm.weight, layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
+        
+        hidden_states = forward_smoothquant_attn(
+            layer.self_attn,
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_cache=kv_caches[i],
+            attn_metadata=attn_metadata,
+            scales=scales,
+            )
+
+        # Fully Connected
+        hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.post_attention_layernorm.weight, layer.post_attention_layernorm.variance_epsilon, layer.mlp.gate_up_proj.smooth_scales)
+        
+        hidden_states = forward_smoothquant_mlp(layer.mlp, hidden_states, scales)
+
+    hidden_states, _ = self.norm(hidden_states, residual)
+    return hidden_states
+

ixformer_sdk/contrib/vllm/layers/mixtral.py

@@ -0,0 +1,331 @@
+import functools
+from typing import Dict, Optional, Tuple
+
+import ixformer.inference.functions as ixf
+import torch
+
+
+def mixtral_decoder_layer_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+    if self.use_int_w8a8:
+        return w8a8_forward(
+            self, positions, hidden_states, kv_cache, attn_metadata, residual
+        )
+    else:
+        return original_forward(
+            self, positions, hidden_states, kv_cache, attn_metadata, residual
+        )
+
+
+def original_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+    # Self Attention
+    if residual is None:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+    else:
+        hidden_states, residual = self.input_layernorm(hidden_states, residual)
+    hidden_states = self.self_attn(
+        positions=positions,
+        hidden_states=hidden_states,
+        kv_cache=kv_cache,
+        attn_metadata=attn_metadata,
+    )
+
+    # Fully Connected
+    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+    hidden_states = self.block_sparse_moe(hidden_states)
+    return hidden_states, residual
+
+
+def dynamic_scaled_int8_quant(x):
+    m, k = x.shape
+    i8_x = x.new_empty([m, k], dtype=torch.int8, device="cuda")
+    i8_scales = torch.empty([m], dtype=torch.float32, device="cuda")
+    ixf.dynamic_scaled_int8_quant(i8_x, x, i8_scales)
+    return i8_x, i8_scales
+
+
+def dynamic_w8a8(x, i8_weight, weight_scale):
+    i8_x, i8_scale = dynamic_scaled_int8_quant(x)
+    m, k = x.shape
+    k, n = i8_weight.shape
+    output = x.new_empty([m, n], dtype=x.dtype, device="cuda")
+    ixf.w8a8(
+        i8_x,
+        i8_weight.transpose(0, 1),
+        i8_scale,
+        weight_scale,
+        output=output,
+        out_dtype=x.dtype,
+    )
+    return output
+
+
+def fused_rms_norm_quant_linear(
+    self,
+    hidden_states,
+    ln_weight,
+    eps,
+    linear_weight,
+    linear_weight_scale,
+    residual=None,
+):
+    # lower rouge
+    # if residual is None:
+    #     residual = hidden_states
+    #     i8_hidden_states, _, i8_scales = ixf.residual_rms_norm_dynamic_int8(
+    #         input=hidden_states,
+    #         weight=ln_weight,
+    #         residual=None,
+    #         eps=eps,
+    #     )
+    # else:
+    #     i8_hidden_states, residual, i8_scales = ixf.residual_rms_norm_dynamic_int8(
+    #         input=hidden_states,
+    #         weight=ln_weight,
+    #         residual=residual,
+    #         eps=eps,
+    #     )
+
+    if residual is None:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+    else:
+        hidden_states, residual = self.input_layernorm(hidden_states, residual)
+    i8_hidden_states, i8_scales = dynamic_scaled_int8_quant(hidden_states)
+
+    qkv = hidden_states.new_empty(hidden_states.shape[0], linear_weight.shape[1])
+    ixf.w8a8(
+        i8_hidden_states,
+        linear_weight.transpose(0, 1),
+        i8_scales,
+        linear_weight_scale,
+        output=qkv,
+        out_dtype=hidden_states.dtype,
+    )
+    return qkv, residual
+
+
+def attention(qkv, positions, kv_cache, attn_metadata, self_attn):
+    q, k, v = qkv.split(
+        [self_attn.q_size, self_attn.kv_size, self_attn.kv_size], dim=-1
+    )
+    q, k = self_attn.rotary_emb(positions, q, k)
+    attn_output = self_attn.attn(q, k, v, kv_cache, attn_metadata)
+    return attn_output
+
+
+def fused_rms_norm_attention(
+    self,
+    hidden_states,
+    ln_weight,
+    eps,
+    positions,
+    kv_cache,
+    attn_metadata,
+    self_attn,
+    residual=None,
+):
+    hidden_states, residual = fused_rms_norm_quant_linear(
+        self,
+        hidden_states,
+        ln_weight,
+        eps,
+        self_attn.qkv_proj.weight,
+        self_attn.qkv_proj.weight_scale,
+        residual,
+    )
+    hidden_states = attention(
+        hidden_states, positions, kv_cache, attn_metadata, self_attn
+    )
+
+    hidden_states = dynamic_w8a8(
+        hidden_states, self_attn.o_proj.weight, self_attn.o_proj.weight_scale
+    )
+    # hidden_states,_ = self_attn.o_proj(hidden_states) # quant+linear+allreduce
+    return hidden_states, residual
+
+
+def w8a8_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+
+    # qkv,_ = self.self_attn.qkv_proj(hidden_states)
+    hidden_states, residual = fused_rms_norm_attention(
+        self,
+        hidden_states,
+        self.input_layernorm.weight,
+        self.input_layernorm.variance_epsilon,
+        positions,
+        kv_cache,
+        attn_metadata,
+        self.self_attn,
+        residual,
+    )
+
+    # allreduce
+    tp_size = self.block_sparse_moe.experts.tp_size
+    if tp_size > 1:
+        from vllm.distributed import tensor_model_parallel_all_reduce
+
+        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+
+    # rms norm
+    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+    # moe
+    hidden_states = fused_moe(
+        hidden_states,
+        self.block_sparse_moe.gate.weight,
+        top_k=self.block_sparse_moe.experts.top_k,
+        w1=self.block_sparse_moe.experts.w13_weight,
+        w2=self.block_sparse_moe.experts.w2_weight,
+        w1_scale=self.block_sparse_moe.experts.w13_weight_scale,
+        w2_scale=self.block_sparse_moe.experts.w2_weight_scale,
+    )
+
+    # allreduce
+    if tp_size > 1:
+        from vllm.distributed import tensor_model_parallel_all_reduce
+
+        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+
+    return hidden_states, residual
+
+
+def fused_experts(hidden_states, router_logits, top_k, w1, w2, w1_scale, w2_scale):
+
+    """
+    Args:
+        hidden_states:          (num_tokens, k)                 dtype
+        router_logits:          (num_tokens, num_experts)       torch.float32
+        top_k                                                   int
+        w1:                     (num_experts, 2n, k)            torch.int8
+        w2:                     (num_experts, k, n)             torch.int8
+        w1_scale:               (num_experts, 2n)               torch.float32
+        w2_scale:               (num_experts, k)                torch.float32
+    Returns
+        final_hidden_states:    (num_tokens, k)                 dtype
+    """
+
+    # topk_weight: (num_tokens, top_k) torch.float32
+    # topk_ids:    (num_tokens, top_k) torch.int32
+    topk_weight, topk_ids = ixf.moe_topk_softmax(
+        gating_output=router_logits,
+        topk=top_k,
+        renormalize=True,
+    )
+
+    dtype = hidden_states.dtype
+    num_tokens, num_experts = router_logits.shape
+    expand_tokens = num_tokens * top_k
+
+    (
+        src_to_dst,
+        sorted_token_ids,
+        expert_sizes_gpu,
+        expert_sizes_cpu,
+    ) = ixf.moe_compute_token_index(
+        topk_ids=topk_ids,
+        num_experts=num_experts,
+    )
+    expert_sizes_cpu = expert_sizes_gpu.cpu()
+
+    # expand + reorder + quant
+    # i8_hidden_states: (expand_tokens, k) torch.int8
+    i8_hidden_states, a_scale = ixf.moe_expand_input_dynamic_scaled_int8(
+        hidden_states=hidden_states,
+        dst_to_src=sorted_token_ids,
+        dst_tokens=expand_tokens,
+        topk=top_k,
+        src_to_dst=src_to_dst,
+        topk_ids=None,  # use smooth quant
+        smooth_scales=None,  # use smooth quant
+    )
+
+    # w8a8 group gemm 1
+    # pt_output_1: (expand_tokens, 2n) dtype
+    pt_output_1 = ixf.moe_w8a8_group_gemm(
+        input=i8_hidden_states,
+        weight=w1,
+        i_scales=a_scale,
+        w_scales=w1_scale,
+        output_dtype=dtype,
+        tokens_per_experts=expert_sizes_cpu,
+        dst_to_src=None,
+        format="TN",
+    )
+
+    # act + quant
+    # pt_output_2: (expand_tokens, n) torch.int8
+    pt_output_2, a2_scale = ixf.activation_dynamic_scaled_int8(
+        input=pt_output_1,
+        bias=None,  # add gemm bias
+        smooth_scales=None,  # use smooth quant
+        dst_to_src=sorted_token_ids,
+        topk_ids=None,  # add gemm bias or use smooth quant
+        act_type="swiglu",
+    )
+
+    # w8a8 group gemm 2 + reorder
+    # pt_output_3: (expand_tokens, k) dtype
+    pt_output_3 = ixf.moe_w8a8_group_gemm(
+        input=pt_output_2,
+        weight=w2,
+        i_scales=a2_scale,
+        w_scales=w2_scale,
+        output_dtype=dtype,
+        tokens_per_experts=expert_sizes_cpu,
+        dst_to_src=sorted_token_ids,
+        format="TN",
+    )
+
+    # mul + reduce_sum
+    # final_hidden_states: (num_tokens, k)
+    final_hidden_states = ixf.moe_output_reduce_sum(
+        input=pt_output_3.view(num_tokens, top_k, -1),
+        topk_weight=topk_weight,
+    )
+
+    return final_hidden_states
+
+
+def fused_moe(hidden_states, gate_weight, top_k, w1, w2, w1_scale, w2_scale):
+    orig_shape = hidden_states.shape
+    hidden_size = hidden_states.shape[-1]
+
+    hidden_states = hidden_states.view(-1, hidden_size)
+
+    # router_logits: (num_tokens, n_experts)
+    # gate_weight: fp16
+    router_logits = ixf.linear(hidden_states, gate_weight)
+    router_logits = router_logits.to(torch.float32)
+
+    final_hidden_states = fused_experts(
+        hidden_states,
+        router_logits,
+        top_k,
+        w1,
+        w2,
+        w1_scale,
+        w2_scale,
+    )
+
+    return final_hidden_states.view(orig_shape)

ixformer_sdk/contrib/vllm/quantize/__init__.py

@@ -0,0 +1,14 @@
+from .smoothquant import smoothquant_prepare_quantize,smoothquant_export_quantized_weights
+from .w8a16 import w8a16_prepare_quantize,w8a16_export_quantized_weights
+
+SUPPORT_METHOD = {
+    "smoothquant": [smoothquant_prepare_quantize,smoothquant_export_quantized_weights],
+    "w8a16": [w8a16_prepare_quantize,w8a16_export_quantized_weights],
+}
+
+def get_quantize_method(method_name:str):
+    try:
+        method = SUPPORT_METHOD[method_name]
+    except:
+        raise ValueError(f"Only support quantization methods: {SUPPORT_METHOD.keys()}, but got {method_name}")
+    return method

ixformer_sdk/contrib/vllm/quantize/smoothquant.py

@@ -0,0 +1,407 @@
+import os
+
+import torch
+
+
+def smoothquant_prepare_quantize(self, quant_params={}):
+    model = self.model_runner.model
+
+    def update_act_scales(act_scales, x):
+        # 动态统计每次输入的最大值
+        hidden_dim = x.shape[-1]
+        x = x.view(-1, hidden_dim).abs().detach()
+        # [k]
+        comming_max = torch.max(x, dim=0, keepdim=True)[0].float()
+
+        if act_scales is None:
+            act_scales = comming_max
+        else:
+            act_scales = torch.max(act_scales, comming_max)
+        return act_scales
+
+    from functools import partial
+
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+
+    def new_forward(input_, m, raw_forward):
+        if not hasattr(m, "act_scales"):
+            m.act_scales = None
+        m.act_scales = update_act_scales(m.act_scales, input_)
+        return raw_forward(input_)
+
+    for name, m in model.named_modules():
+        if (
+            isinstance(m, QKVParallelLinear)
+            or isinstance(m, RowParallelLinear)
+            or isinstance(m, MergedColumnParallelLinear)
+            or isinstance(m, ColumnParallelLinear)
+        ):
+            m.forward = partial(new_forward, m=m, raw_forward=m.forward)
+
+def smoothquant_export_quantized_weights(self, save_path, quant_params={}):
+    gb_per_file = quant_params.get("filesize_limit", None)
+    smooth_alpha = quant_params.get("smooth_alpha", 0.5)
+    dynamic_quant_type = quant_params.get("dynamic_quant_type", "gpu")
+    assert dynamic_quant_type in ["gpu","cpu","kernel"]
+    if self.rank == 0:
+        print(f"set smooth_alpha={smooth_alpha}")
+        print(f"use quantize weight type: {dynamic_quant_type}")
+    
+    import ixformer._C as ops
+    def per_token_quant_8bit(weight):
+        # weight: [m,k]
+        dtype = weight.dtype
+        i8_weight = weight
+        scale = i8_weight.abs().max(dim=-1, keepdim=True)[0] / 127
+        i8_weight = i8_weight / scale.to(dtype)
+        i8_weight = torch.clamp(torch.round(i8_weight), -128, 127).to(torch.int8)
+        return i8_weight, scale.float()
+
+    def smooth_quant_weight_gpu_cpu(weight, act_scale, alpha=0.5, device="cpu"):
+        device = torch.device("cpu") if device == "cpu" else weight.device
+        ori_dtype = weight.dtype
+        # [1, k]
+        act_scale = act_scale.float().to(device).view(1, -1)
+        weight = weight.to(device)
+        # [1, k]
+        weight_scale = weight.abs().max(dim=0, keepdim=True)[0].float()
+        if alpha == -1:
+            smooth_scales = torch.ones_like(act_scale)
+        else:
+            smooth_scales = act_scale.pow(alpha) / weight_scale.pow(1 - alpha).clamp(
+                min=1e-5
+            )
+        weight = weight * smooth_scales.to(ori_dtype)
+        i8_weight, weight_scales = per_token_quant_8bit(weight)
+        # 为了可以使用 input * smooth_scales
+        if alpha == -1:
+            smooth_scales = torch.ones_like(act_scale)
+        else:
+            smooth_scales = weight_scale.pow(1 - alpha) / act_scale.pow(alpha).clamp(
+                min=1e-5
+            )
+        return i8_weight, weight_scales, smooth_scales.to(ori_dtype)
+    
+    def smooth_quant_weight_kernel(weight, act_scale, alpha=0.5):
+        output = torch.zeros_like(weight,dtype=torch.int8)
+        weight_scales = torch.zeros(weight.shape[:-1],dtype=torch.float, device=weight.device)
+        weight_max = torch.zeros(weight.shape[-1],dtype=torch.float, device=weight.device)
+        smooth_scales = torch.zeros(weight.shape[-1],dtype=weight.dtype, device=weight.device)
+        ops.infer.weight_quant_smoothquant(
+            weight, act_scale, alpha, output, weight_scales, smooth_scales, weight_max
+        )
+        return output, weight_scales.view(-1,1), smooth_scales.view(1,-1)
+    
+    def smooth_quant_weight(weight, act_scale, alpha=0.5):
+        if dynamic_quant_type == "kernel":
+            return smooth_quant_weight_kernel(weight,act_scale,alpha)
+        else:
+            return smooth_quant_weight_gpu_cpu(weight,act_scale,alpha,dynamic_quant_type)
+
+    model = self.model_runner.model
+
+    from vllm.distributed import (
+        tensor_model_parallel_all_gather,
+        tensor_model_parallel_all_reduce,
+        get_tensor_model_parallel_world_size
+    )
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.model_executor.layers.vocab_parallel_embedding import (
+        ParallelLMHead,
+        VocabParallelEmbedding,
+    )
+    from vllm.model_executor.models.falcon import FalconForCausalLM
+
+    for name, m in model.named_modules():
+        if isinstance(m, VocabParallelEmbedding):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+            
+        elif isinstance(m, ParallelLMHead):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            # bias shape: [vocab_size // tp]
+            if m.bias is not None:
+                bias = tensor_model_parallel_all_gather(m.bias, dim=0)
+                bias = bias[:m.org_vocab_size].contiguous()
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, QKVParallelLinear):
+            # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
+            # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
+            if self.parallel_config.world_size > 1:
+                total_q_hidden_size = m.total_num_heads * m.head_size
+                partial_q_hidden_size = m.num_heads * m.head_size
+                total_kv_hidden_size = m.total_num_kv_heads * m.head_size
+                partial_kv_hidden_size = m.num_kv_heads * m.head_size
+                
+                if m.bias is not None:
+                    # TODO do not support padding..
+                    bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
+                    q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+                    k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    
+                    q_bias[:] = m.bias[:partial_q_hidden_size]
+                    k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
+                    v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
+                    
+                    bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
+
+                q_tensor = m.weight.new_zeros(m.total_num_heads * m.head_size, m.weight.shape[1])
+                k_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
+                v_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
+                
+                q_in_weight = m.weight[:-m.num_kv_heads * m.head_size * 2]
+                k_in_weight = m.weight[-m.num_kv_heads * m.head_size * 2:-m.num_kv_heads * m.head_size]
+                v_in_weight = m.weight[-m.num_kv_heads * m.head_size:]
+                
+                if getattr(m,"start_idx",None) is not None:
+                    start_idx = getattr(m,"start_idx")
+                    weight_end_idx = m.num_heads * m.head_size if not getattr(m,"is_padding") else (m.num_heads - 1) * m.head_size
+                    end_idx = start_idx + weight_end_idx
+                else:
+                    start_idx = self.rank * m.num_heads * m.head_size
+                    weight_end_idx = m.num_heads * m.head_size
+                    end_idx = start_idx + weight_end_idx
+                assert q_tensor[start_idx:end_idx,:].shape == q_in_weight[:weight_end_idx, :].shape
+                q_tensor[start_idx:end_idx,:] = q_in_weight[:weight_end_idx, :]
+                
+                if m.num_kv_head_replicas > 1:
+                    if self.rank % m.num_kv_head_replicas == 0:
+                        rank = self.rank // m.num_kv_head_replicas
+                        k_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
+                        v_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
+                else:
+                    k_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
+                    v_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
+                
+                q_tensor = tensor_model_parallel_all_reduce(q_tensor)
+                k_tensor = tensor_model_parallel_all_reduce(k_tensor)
+                v_tensor = tensor_model_parallel_all_reduce(v_tensor)
+                
+                if isinstance(model, FalconForCausalLM):
+                    num_query_heads_per_kv_head = (
+                        m.total_num_heads // m.total_num_kv_heads
+                    )
+                    q_tensor = q_tensor.view(
+                        m.total_num_kv_heads,
+                        num_query_heads_per_kv_head,
+                        m.head_size,
+                        -1,
+                    )
+                    k_tensor = k_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
+                    v_tensor = v_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
+                    weight_tensor = torch.cat(
+                        [q_tensor, k_tensor, v_tensor], dim=1
+                    ).view(-1, m.hidden_size)
+                else:
+                    weight_tensor = torch.cat([q_tensor, k_tensor, v_tensor])
+                assert (
+                    weight_tensor.shape[0]
+                    == total_q_hidden_size + total_kv_hidden_size * 2
+                )
+                assert weight_tensor.shape[1] == m.hidden_size
+
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+            
+        elif isinstance(m, MergedColumnParallelLinear):
+            if self.parallel_config.world_size > 1:
+                # weight shape: [intermediate_size // tp * 2, hidden_size]
+                # bias shape: [intermediate_size // tp * 2]
+                output_sizes = m.output_sizes
+                output_size = sum(output_sizes)
+                partial_output_sizes = [
+                    i // self.parallel_config.world_size for i in output_sizes
+                ]
+
+                if m.bias is not None:
+                    index_start = 0
+                    partial_index_start = 0
+                    bias_tenosr = m.bias.new_zeros(output_size)
+                    for i in range(len(output_sizes)):
+                        index_out = index_start + output_sizes[i]
+                        sub_bias_tensor = bias_tenosr[index_start:index_out]
+                        partial_size = partial_output_sizes[i]
+                        sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
+                        
+                        index_start += output_sizes[i]
+                        partial_index_start += partial_size
+                    bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
+                
+                weight_tensor = m.weight.new_zeros(output_size, m.input_size)
+
+                idx_out_start = 0
+                idx_partial_satrt = 0
+                for i in range(len(output_sizes)):
+                    idx_out_end = idx_out_start + output_sizes[i]
+                    sub_weight_tensor = weight_tensor[idx_out_start:idx_out_end]
+                    partial_size = partial_output_sizes[i]
+                    sub_weight_tensor[
+                        self.rank * partial_size : (self.rank + 1) * partial_size
+                    ] = m.weight[idx_partial_satrt : idx_partial_satrt + partial_size]
+
+                    idx_out_start += output_sizes[i]
+                    idx_partial_satrt += partial_size
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, ColumnParallelLinear):
+            # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
+            # bias shape: [some_dim // tp]
+            if m.bias is not None:
+                bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
+
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                   weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, RowParallelLinear):
+            # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
+            # bias shape: [hidden_size]
+            if m.bias is not None:
+                bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if getattr(m,"start_idx", None) is not None:
+                start_idx = getattr(m,"start_idx")
+                end_idx = start_idx + (m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size))
+                weight_end_idx = m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size)
+            else:
+                start_idx = m.input_size_per_partition * self.rank
+                end_idx = start_idx + m.input_size_per_partition
+                weight_end_idx = m.input_size_per_partition
+            
+            act_scales = m.act_scales.new_zeros(m.input_size)
+            assert act_scales[start_idx:end_idx].shape == m.act_scales.view(-1)[:weight_end_idx].shape
+            act_scales[start_idx:end_idx] = m.act_scales.view(-1)[:weight_end_idx]
+            act_scales = tensor_model_parallel_all_reduce(act_scales)
+            m.act_scales = act_scales
+            
+            weight_tensor = m.weight.new_zeros(m.weight.shape[0],m.input_size)
+            assert  weight_tensor[:,start_idx:end_idx].shape == m.weight[:,:weight_end_idx].shape
+            weight_tensor[:,start_idx:end_idx] = m.weight[:,:weight_end_idx]
+            weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                smooth_scales = smooth_scales.view(1,-1)
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+        else:
+            pass
+        
+        torch.cuda.empty_cache()
+
+    # save weights
+    if self.is_driver_worker:
+        from safetensors.torch import save_file
+        
+        tensors = {}
+        saved = False
+        count = 0
+        size_in_bytes = 0
+        
+        tensors = {}
+        for name, weight in model.named_parameters():
+            if "act_scales" in name:
+                continue
+            # skip lm_head_weight if needed..
+            if "lm_head" in name and model.config.tie_word_embeddings:
+                continue
+            tensors[name] = weight
+            
+            saved = False
+            if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
+                weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
+                save_file(tensors, weight_path)
+                print(f"The quantified weights were successfully saved in {weight_path}.")
+                tensors.clear()
+                saved = True
+                count += 1
+                size_in_bytes = 0
+        
+        if not saved: 
+            weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))   
+            save_file(tensors, weight_path)
+            print(f"The quantified weights were successfully saved in {weight_path}.")

ixformer_sdk/contrib/vllm/quantize/w8a16.py

@@ -0,0 +1,233 @@
+import os
+
+import torch
+
+def w8a16_prepare_quantize(self, quant_params={}):
+    # We need do nothing in here
+    pass
+
+
+def w8a16_export_quantized_weights(self, save_path, quant_params={}):
+    gb_per_file = quant_params.get("filesize_limit", None)
+    int8_min = -127
+
+    def w8a16_quantization(weight):
+        # all weights should be [output,input], otherwise, we may get an wrong weight and scale...
+        scale = torch.abs(weight).max(dim=-1)[0] / 127.0
+        int8_weight = torch.clamp(weight / scale.view(-1,1),min=int8_min,max=127).to(torch.int8).contiguous()
+        scale = scale.view(1,-1).contiguous()
+        return int8_weight, scale
+
+
+    model = self.model_runner.model
+
+    from vllm.distributed.communication_op import (
+        tensor_model_parallel_all_gather,
+        tensor_model_parallel_all_reduce,
+    )
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.model_executor.layers.vocab_parallel_embedding import (
+        ParallelLMHead,
+        VocabParallelEmbedding,
+    )
+
+    for name, m in model.named_modules():
+        if isinstance(m, VocabParallelEmbedding):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, ParallelLMHead):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            # bias shape: [vocab_size // tp]
+            if m.bias is not None:
+                bias = tensor_model_parallel_all_gather(m.bias, dim=0)
+                bias = bias[:m.org_vocab_size].contiguous()
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, QKVParallelLinear):
+            # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
+            # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
+            if self.parallel_config.world_size > 1:
+                total_q_hidden_size = m.total_num_heads * m.head_size
+                partial_q_hidden_size = m.num_heads * m.head_size
+                total_kv_hidden_size = m.total_num_kv_heads * m.head_size
+                partial_kv_hidden_size = m.num_kv_heads * m.head_size
+
+                if m.bias is not None:
+                    bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
+                    q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+                    k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    
+                    q_bias[:] = m.bias[:partial_q_hidden_size]
+                    k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
+                    v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
+                    
+                    bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
+                
+                weight_tensor = m.weight.new_zeros(
+                    total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size
+                )
+
+                q_tensor = weight_tensor[:total_q_hidden_size, :]
+                q_tensor = q_tensor[self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+
+                k_tensor = weight_tensor[total_q_hidden_size : total_q_hidden_size + total_kv_hidden_size]
+                k_tensor = k_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+
+                v_tensor = weight_tensor[total_q_hidden_size + total_kv_hidden_size :]
+                v_tensor = v_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+
+                q_tensor[:, :] = m.weight[: partial_q_hidden_size, :]
+                k_tensor[:, :] = m.weight[partial_q_hidden_size : partial_q_hidden_size + partial_kv_hidden_size, :]
+                v_tensor[:, :] = m.weight[partial_q_hidden_size + partial_kv_hidden_size : , :]
+
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+                    
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, MergedColumnParallelLinear):
+            # weight shape: [intermediate_size // tp * 2, hidden_size]
+            # bias shape: [intermediate_size // tp * 2]
+            if self.parallel_config.world_size > 1:
+                output_sizes = m.output_sizes
+                output_size = sum(output_sizes)
+                partial_output_sizes = [
+                    i // self.parallel_config.world_size for i in output_sizes
+                ]
+                
+                if m.bias is not None:
+                    index_start = 0
+                    partial_index_start = 0
+                    bias_tenosr = m.bias.new_zeros(output_size)
+                    for i in range(len(output_sizes)):
+                        index_out = index_start + output_sizes[i]
+                        sub_bias_tensor = bias_tenosr[index_start:index_out]
+                        partial_size = partial_output_sizes[i]
+                        sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
+                        
+                        index_start += output_sizes[i]
+                        partial_index_start += partial_size
+                    bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
+
+                weight_tensor = m.weight.new_zeros(output_size, m.input_size)
+
+                index_start = 0
+                partial_index_start = 0
+                for i in range(len(output_sizes)):
+                    index_out = index_start + output_sizes[i]
+                    sub_weight_tensor = weight_tensor[index_start:index_out]
+                    partial_size = partial_output_sizes[i]
+                    sub_weight_tensor[self.rank * partial_size : (self.rank + 1) * partial_size] = m.weight[partial_index_start : partial_index_start + partial_size]
+
+                    index_start += output_sizes[i]
+                    partial_index_start += partial_size
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+                
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        
+        elif isinstance(m, ColumnParallelLinear):
+            # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
+            # bias shape: [some_dim // tp]
+            if m.bias is not None:
+                bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        
+        elif isinstance(m, RowParallelLinear):
+            # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
+            # bias shape: [hidden_size]
+            if m.bias is not None:
+                bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=-1)
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        else:
+            pass
+        
+        torch.cuda.empty_cache()
+    
+    # save weights
+    if self.is_driver_worker:
+        from safetensors.torch import save_file
+
+        tensors = {}
+        saved = False
+        count = 0
+        size_in_bytes = 0
+        
+        for name, weight in model.named_parameters():
+            if "lm_head" in name and model.config.tie_word_embeddings:
+                continue
+            size_in_bytes += weight.numel() * weight.element_size()
+            tensors[name] = weight
+            saved = False
+            if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
+                weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
+                save_file(tensors, weight_path)
+                print(f"The quantified weights were successfully saved in {weight_path}.")
+                tensors.clear()
+                saved = True
+                count += 1
+                size_in_bytes = 0
+        
+        if not saved: 
+            weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))   
+            save_file(tensors, weight_path)
+            print(f"The quantified weights were successfully saved in {weight_path}.")

ixformer_sdk/contrib/vllm_flash_attn/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "2.6.1"
+
+from .flash_attn_interface import *

ixformer_sdk/core/config.py

@@ -0,0 +1,184 @@
+import os
+from typing import Callable, Optional
+
+# =========================================================
+# Utils
+# =========================================================
+
+
+def number_type(scalar_type):
+    def wrap(val: Optional[str]):
+        if val is None:
+            return None
+
+        return scalar_type(val)
+
+    return wrap
+
+
+def bool_type(val: Optional[str]):
+    if val is None:
+        return False
+
+    if isinstance(val, str):
+        return val.lower() in ["1", "t", "true"]
+
+    if isinstance(val, int):
+        return val != 0
+
+    raise RuntimeError(f"Invalid bool type, got {type(val), val}")
+
+
+def list_type(scalar_type=str):
+    def wrap(val: Optional[str]):
+        if val is None:
+            return []
+
+        if not isinstance(val, str):
+            raise RuntimeError(
+                f"list_type: Got invalid type, expect str, but got {val}."
+            )
+
+        return [scalar_type(v) for v in val.split(",")]
+
+    return wrap
+
+
+def Field(
+    name: str,
+    static: bool = True,
+    type: Callable = str,
+    choices: Optional[list] = None,
+    help: Optional[str] = None,
+    **kwargs,
+):
+    """
+    Define environment variable field
+
+    Example:
+    Static mode:
+        # define
+        ENABLE_XX = Field("ENABLE_XX", type=bool, help="ENABLE_XX")
+
+        # use
+        config.ENABLE_XX
+
+    Dynamic mode:
+        # Please use lowercase naming to differentiate it with static mode.
+
+        # define
+        enable_cc = Field("ENABLE_CC", type=bool, static=False, help="enable_cc")
+
+        # use
+        config.enable_cc()
+
+    Set default value:
+        # define
+        ENABLE_TT = Field("ENABLE_TT", type=bool, default=False, help="ENABLE_TT")
+
+        # use
+        config.ENABLE_TT
+
+    Use list:
+        # define
+        CUDA_VISIBLE_DEVICES = Field("CUDA_VISIBLE_DEVICES", type=list_type(int), help="CUDA_VISIBLE_DEVICES")
+
+        # use
+        # the CUDA_VISIBLE_DEVICES is parsed to list, and it's value is int type.
+        for device_id in CUDA_VISIBLE_DEVICES:
+            ...
+
+    """
+
+    if type == bool:
+        type = bool_type
+
+    elif type in [list, tuple]:
+        type = list_type(scalar_type=str)
+
+    elif type in [int, float]:
+        type = number_type(type)
+
+    if static:
+        env_val = type(os.environ.get(name, **kwargs))
+        if choices is not None and env_val is not None and env_val not in choices:
+            raise RuntimeError(
+                f"Got invalid value, expect {choices}, but got {env_val}."
+            )
+        return env_val
+
+    def _get():
+        env_val = type(os.environ.get(name, **kwargs))
+        if choices is not None and env_val is not None and env_val not in choices:
+            raise RuntimeError(
+                f"Got invalid value, expect {choices}, but got {env_val}."
+            )
+        return env_val
+
+    return _get
+
+
+# =========================================================
+# Functions Config
+# =========================================================
+
+IXFORMER_GEMV_THRESHOLD = Field(
+    "IXFORMER_GEMV_THRESHOLD",
+    type=int,
+    default=1,
+    help="Set the threshold for using gemv.",
+)
+
+
+# =========================================================
+# Distributed Config
+# =========================================================
+
+IXFORMER_COMM_SHM_SIZE = Field(
+    "IXFORMER_COMM_SHM_SIZE",
+    type=int,
+    default=None,
+    help="set shared memory size of ipc comm.",
+)
+
+IXFORMER_ENABLE_OVERLAP_COMM = Field(
+    "IXFORMER_ENABLE_OVERLAP_COMM",
+    type=bool,
+    default=False,
+    help="enable overlap communcation and compute.",
+)
+
+IXFORMER_OVERLAP_GEMM_METHOD = Field(
+    "IXFORMER_OVERLAP_GEMM_METHOD",
+    type=int,
+    default=None,
+    choices=[0, 1],
+    help="set gemm backend, 0: ixinfer, 1: cublas.",
+)
+
+IXFORMER_OVERLAP_CHUNKS = Field(
+    "IXFORMER_OVERLAP_CHUNKS", type=int, default=2, help="set split chunks."
+)
+
+IXFORMER_OVERLAP_SPLIT_RATIO = Field(
+    "IXFORMER_OVERLAP_SPLIT_RATIO",
+    type=float,
+    default=None,
+    help="set split chunks ratio.",
+)
+
+IXFORMER_PAGED_ATTENTION_ALGO = Field(
+    "IXFORMER_PAGED_ATTENTION_ALGO",
+    type=str,
+    default="ixinfer",
+    choices=["ixinfer", "ixformer"],
+    help="set paged attention algo.",
+)
+
+IXFORMER_UNPAD_ATTENTION_ALGO = Field(
+    "IXFORMER_UNPAD_ATTENTION_ALGO",
+    type=str,
+    default="ixinfer",
+    choices=["ixinfer", "ixinfer-ex"],
+    help="set enpad attention algo.",
+)

ixformer_sdk/core/dispatcher.py

@@ -0,0 +1,20 @@
+class Dispatcher(object):
+    """
+    create object by dispatcher to reuse object.
+    """
+
+    _dispatcher = dict()
+
+    @classmethod
+    def dispatcher(cls, *args, **kwargs):
+        key = cls.dispatcher_key(*args, **kwargs)
+        obj = cls._dispatcher.get(key, None)
+        if obj is None:
+            obj = cls(*args, **kwargs)
+            cls._dispatcher[key] = obj
+
+        return obj
+
+    @classmethod
+    def dispatcher_key(cls, *args, **kwargs):
+        raise NotImplementedError()

ixformer_sdk/core/multi_level_cache.py

@@ -0,0 +1,54 @@
+class MultiLevelCache(object):
+    def __init__(self):
+        self._l1_key = None
+        self._l1_value = None
+
+        self._l2_size = 3
+        self._l2 = [(None, None) for _ in range(self._l2_size)]
+        self._l2_ptr = 0
+
+        self._l3 = dict()
+
+    def set(self, key, value):
+        self._l1_key = key
+        self._l1_value = value
+
+        self._l2[self._l2_ptr] = (key, value)
+        self._l2_ptr = (self._l2_ptr + 1) % 3  # l2_size: 3
+
+        self._l3[key] = value
+
+    def get(self, key, *args):
+        if key == self._l1_key:
+            return self._l1_value
+
+        l2 = self._l2
+        if key == l2[0][0]:
+            return l2[0][1]
+
+        if key == l2[1][0]:
+            return l2[1][1]
+
+        if key == l2[2][0]:
+            return l2[2][1]
+
+        return self._l3.get(key, *args)
+
+    def containe(self, key):
+        return key in self._l3
+
+    def __getitem__(self, item):
+        return self.get(item)
+
+    def __setitem__(self, key, value):
+        self.set(key, value)
+
+    def __contains__(self, item):
+        if item == self._l1_key:
+            return True
+
+        l2 = self._l2
+        if item == l2[0][0] or item == l2[1][0] or item == l2[2][0]:
+            return True
+
+        return item in self._l3

ixformer_sdk/core/operator_autotuning.py

@@ -0,0 +1,237 @@
+import abc
+import bisect
+import functools
+import itertools
+import random
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+
+import ixformer.distributed as ixfd
+from ixformer.utils.benchmark.cuda_benchmark import Functor, cuda_benchmark
+
+
+def sync_ranks_metric(value, group=None):
+    if not isinstance(value, (torch.Tensor, int, float)):
+        raise RuntimeError(
+            f"Invalid metric value, expect `Tensor`, `int`, or `float` type, but got {value}."
+        )
+
+    if torch.is_tensor(value):
+        value = value.to("cuda")
+    else:
+        value = torch.tensor([value], dtype=torch.float, device="cuda")
+
+    dist.broadcast(value, src=0, group=group)
+    return value.cpu().item()
+
+
+class AutotuningFinder(object):
+    def freeze(self):
+        pass
+
+    @abc.abstractmethod
+    def get(self, key) -> Callable:
+        pass
+
+    @abc.abstractmethod
+    def set(self, *args, **kwargs):
+        pass
+
+
+class BasedKeyFinder(AutotuningFinder):
+    def __init__(self):
+        self._key_to_value: Dict[Any, Callable] = dict()
+
+    def get(self, key, **kwargs) -> Callable:
+        if "default" in kwargs:
+            return self._key_to_value.get(key, kwargs["default"])
+        return self._key_to_value[key]
+
+    def set(self, key, value):
+        self._key_to_value[key] = value
+
+    def containe(self, key):
+        return key in self._key_to_value
+
+
+class TreeNode:
+    def __init__(self):
+        self.nodes: List[Union[Any, TreeNode]] = list()
+        self.key_to_nodes: Dict[Any, TreeNode] = dict()
+
+    def add(self, key, value):
+        if isinstance(key, (tuple, list)):
+            if len(key) == 1:
+                self.insert_value(key[0], value)
+            else:
+                self.recurse_add_node(key, value)
+        else:
+            self.insert_value(key, value)
+
+    def insert_value(self, key, value):
+        self.nodes.append((key, value))
+        self.key_to_nodes[key] = value
+
+    def recurse_add_node(self, key, value):
+        if key[0] in self.key_to_nodes:
+            node = self.key_to_nodes[key[0]]
+        else:
+            node = TreeNode()
+            self.key_to_nodes[key[0]] = node
+            self.insert_value(key[0], node)
+
+        node.add(key[1:], value)
+
+    def sort(self):
+        self.nodes.sort(key=lambda x: x[0])
+        for _, node in self.nodes:
+            if isinstance(node, TreeNode):
+                node.sort()
+
+    def find(self, key):
+        is_list_key = isinstance(key, (tuple, list))
+        if not is_list_key:
+            key = (key,)
+
+        num_querys = len(key)
+        node = self
+        for key_idx in range(num_querys):
+            query_key = key[key_idx]
+            idx = bisect.bisect_left(node.nodes, (query_key,)) - 1
+            if idx <= 0:
+                node = node.nodes[0][1]
+            elif idx >= len(node.nodes):
+                node = node.nodes[-1][1]
+            else:
+                node = node.nodes[idx][1]
+
+        return node
+
+    def show(self, indent=0):
+        for k, node in self.nodes:
+            print(" " * indent, end="")
+            if isinstance(node, TreeNode):
+                print(f"key: {k}")
+                node.show(indent=indent + 4)
+            else:
+                print(f"key: {k}, node: {node}")
+
+
+class BasedRangeFinder(AutotuningFinder):
+    def __init__(self):
+        super().__init__()
+
+        self.tree = TreeNode()
+        self._found_cache: Dict[Any, Callable] = dict()
+
+    def freeze(self):
+        self.tree.sort()
+
+    def get(self, key) -> Callable:
+        value = self._found_cache.get(key, None)
+        if value is not None:
+            return value
+
+        value = self.tree.find(key)
+        self._found_cache[key] = value
+        return value
+
+    def set(self, key, value):
+        self.tree.add(key, value)
+
+
+class OperatorAutotuning(object):
+    def __init__(self, num_repeated=5, num_warmup=3, dist_barrier=False):
+        self.num_repeated = num_repeated
+        self.num_warmup = num_warmup
+        self.dist_barrier = dist_barrier
+
+    @abc.abstractmethod
+    def operators(self):
+        raise NotImplementedError()
+
+    def __call__(self, *args, **kwargs):
+        return self.exec_best_operator(args, kwargs)
+
+    @abc.abstractmethod
+    def exec_best_operator(self, args, kwargs):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def autotuning(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def perf_best_operator(self, *args, **kwargs) -> Callable:
+        best_operator = None
+        best_operator_time = float("inf")
+
+        for idx, operator in enumerate(self.operators()):
+            op_time = self.perf_operator_time(operator, *args, **kwargs)
+            if op_time < best_operator_time:
+                best_operator = operator
+                best_operator_time = op_time
+
+            # print(operator, op_time)
+
+        return best_operator
+
+    def perf_operator_time(self, op: Callable, *args, **kwargs) -> float:
+        fn = Functor(op, *args, **kwargs)
+        time = cuda_benchmark(fn, self.num_repeated, self.num_warmup, self.dist_barrier)
+        return time.gpu
+
+
+class OperatorRuntimeAutotuning(OperatorAutotuning):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.operator_finder = BasedKeyFinder()
+
+    @abc.abstractmethod
+    def get_operator_key(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def exec_best_operator(self, args, kwargs):
+        key = self.get_operator_key(*args, **kwargs)
+        operator = self.operator_finder.get(key, default=None)
+        if operator is None:
+            operator = self.perf_best_operator(*args, **kwargs)
+            self.operator_finder.set(key, operator)
+
+        return operator(*args, **kwargs)
+
+
+class OperatorPreBaseRangeAutotuning(OperatorAutotuning):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.operator_finder = BasedRangeFinder()
+        self._finished_autotuning = False
+
+    @abc.abstractmethod
+    def get_operator_key(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def generate_operator_inputs(self) -> Iterable[Tuple[Tuple, Dict]]:
+        raise NotImplementedError()
+
+    def exec_best_operator(self, args, kwargs):
+        if not self._finished_autotuning:
+            self.autotuning()
+
+        best_op = self.operator_finder.get(self.get_operator_key(*args, **kwargs))
+        return best_op(*args, **kwargs)
+
+    def autotuning(self):
+        for op_args, op_kwargs in self.generate_operator_inputs():
+            best_op = self.perf_best_operator(*op_args, **op_kwargs)
+            self.operator_finder.set(
+                self.get_operator_key(*op_args, **op_kwargs), best_op
+            )
+
+        self.operator_finder.freeze()
+        self._finished_autotuning = True
+        # self.operator_finder.tree.show()

ixformer_sdk/csrc/FindIXFORMER.cmake

@@ -0,0 +1,40 @@
+# use python to find ixformer libs and include
+if (CMAKE_VERSION VERSION_LESS 3.18)
+    set(DEV_MODULE Development)
+else()
+    set(DEV_MODULE Development.Module)
+endif()
+
+find_package(Python COMPONENTS Interpreter ${DEV_MODULE} REQUIRED)
+
+# find ixformer
+set(IXFORMER_FOUND FALSE)
+
+if("${Python_FOUND}" STREQUAL "TRUE")
+    execute_process(
+            COMMAND ${Python_EXECUTABLE} -c "import os, ixformer; print(os.path.dirname(ixformer.__file__))"
+            OUTPUT_VARIABLE IXFORMER_PYDIR
+            ERROR_VARIABLE PYTHON_ERROR
+            RESULT_VARIABLE PYTHON_RESULT
+            OUTPUT_STRIP_TRAILING_WHITESPACE
+            ERROR_STRIP_TRAILING_WHITESPACE
+    )
+    if ("${IXFORMER_PYDIR}" STREQUAL "")
+        message("-- Not found ixFormer")
+    else ()
+        message("-- Found ixFormer: ${IXFORMER_PYDIR}")
+        set(IXFORMER_FOUND TRUE)
+    endif ()
+endif()
+
+set(IXFORMER_COMM_LIBS "ixformer_comm")
+set(IXFORMER_KERNEL_LIBS "ixformer_kernels")
+set(IXFORMER_LIBS "${IXFORMER_COMM_LIBS} ${IXFORMER_KERNEL_LIBS}")
+set(IXFORMER_INCLUDE "")
+set(IXFORMER_DIR "")
+
+if("${IXFORMER_FOUND}" STREQUAL "TRUE")
+    set(IXFORMER_INCLUDE "${IXFORMER_PYDIR}/csrc/include")
+    set(IXFORMER_DIR "${IXFORMER_PYDIR}")
+    message("-- ixFormer LIBS: ${IXFORMER_LIBS}, INCLUDE: ${IXFORMER_INCLUDE}")
+endif ()

ixformer_sdk/csrc/include/ixformer/comm/ccl.h

@@ -0,0 +1,357 @@
+#pragma once
+
+#include "core/op_algo.h"
+#include "nccl.h"
+
+namespace ixformer::comm {
+
+const uint8_t MAX_TENSOR_NDIM = 8;
+constexpr size_t DEFAULT_SHM_SIZE = 16 * 1024 * 1024 * sizeof(float);
+
+struct Comm;
+typedef Comm *Comm_t;
+
+
+struct TensorDesc {
+    void *data_ptr;
+    ncclDataType_t dtype;
+    uint64_t numel;
+    uint8_t ndim;
+    int64_t shape[MAX_TENSOR_NDIM];
+    int64_t stride[MAX_TENSOR_NDIM];
+    bool contiguous;
+};
+
+
+/**
+ * @brief Generate unique communicator id.
+ *
+ * Generates an Id to be used in ncclCommInitRank. ncclGetUniqueId should be
+ * called once and the Id should be distributed to all ranks in the
+ * communicator before calling ncclCommInitRank.
+ *
+ * @param commId: the unique id of communicator, it is created in main rank, and broadcast other rank.
+ */
+void getUniqueId(ncclUniqueId *commId);
+
+/**
+ * @brief Serialize commId to string.
+ * @param commId: the unique id of communicator.
+ * @return: serialized string.
+ */
+std::string serializeUniqueId(const ncclUniqueId &commId);
+
+/**
+ * @brief Deserialize the string of commId.
+ * @param commIdStr: serialized string by serializeUniqueId.
+ * @param commId: output commId
+ */
+void deserializeUniqueId(const std::string &commIdStr, ncclUniqueId *commId);
+
+/**
+ * @brief Creates a new communicator (multi process version).
+ *
+ * Rank must be between 0 and nranks-1 and unique within a communicator clique.
+ * Each rank is associated to a CUDA device, which has to be set before calling ncclCommInitRank.
+ *
+ * It is important to ensure that the current process's CUDA device is set by cudaSetDevice before calling this function,
+ * otherwise, an exception will be thrown.
+ *
+ * @param comm: Communicator
+ * @param nranks: the number of ranks.
+ * @param commId: the unique id of communicator.
+ * @param rank: the rank of current process
+ * @param shm_size: Unlike NCCL, IxFormer communication relies on CUDA IPC for communication by shared memory.
+ *        If the shm_size is not provided, it will use the default value: DEFAULT_SHM_SIZE.
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void initRank(Comm_t *comm, int nranks, ncclUniqueId commId, int rank, size_t shm_size = DEFAULT_SHM_SIZE);
+
+/**
+ * @brief Finalize a communicator.
+ *
+ * ncclCommFinalize flushes all issued communications,
+ * and marks communicator state as ncclInProgress. The state will change to ncclSuccess
+ * when the communicator is globally quiescent and related resources are freed; then,
+ * calling ncclCommDestroy can locally free the rest of the resources (e.g. communicator
+ * itself) without blocking.
+ *
+ * @param comm: Communicator
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void destroy(Comm_t comm);
+
+void delete_comm_resuouces(Comm_t comm);
+
+/**
+ * @brief Whether is initiated.
+ * @param comm: Communicator
+ */
+bool isInitiated(Comm_t comm);
+
+/**
+ * @brief Gets ncclComm_t.
+ * @param comm: Communicator
+ */
+ncclComm_t getNcclComm(Comm_t comm);
+
+/**
+ * @brief Gets the number of ranks in the communicator clique
+ * @param comm: Communicator
+ */
+int getWorldSize(Comm_t comm);
+
+/**
+ * @brief Gets the number of nodes in the communicator clique
+ * @param comm: Communicator
+ */
+int getNumNodes(Comm_t comm);
+
+/**
+ * @brief Returns the user-ordered "rank" associated with the communicator.
+ * @param comm: Communicator
+ */
+int getRank(Comm_t comm);
+
+/**
+ * @brief Returns the cuda device number associated with the communicator.
+ * @param comm: Communicator
+ */
+int getDevice(Comm_t comm);
+
+/**
+ * @brief Gets shared memory size in the communicator clique
+ * @param comm: Communicator
+ */
+uint64_t getIpcShmSize(Comm_t comm);
+
+
+// ============================================================================
+// Collective communication operations
+//
+// Collective communication operations must be called separately for each
+// communicator in a communicator clique.
+//
+// They return when operations have been enqueued on the CUDA stream.
+//
+// Since they may perform inter-CPU synchronization, each call has to be done
+// from a different thread or process, or need to use Group Semantics (see
+// below).
+// ============================================================================
+
+/**
+ * @brief Barrier the member of the communicator
+ * @param comm: Communicator
+ * @param stream: CUDA Stream
+ */
+void barrier(Comm_t comm, cudaStream_t stream);
+
+/**
+ * @brief All-Gather
+ *
+ * Each device gathers sendcount values from other GPUs into senddata,
+ * receiving data from rank i at offset i*sendcount.
+ * Assumes recvcount is equal to nranks*sendcount, which means that recvdata
+ * should have a size of at least nranks*sendcount elements.
+ *
+ * In-place operations will happen if senddata == recvdata + rank * sendcount.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param sendcount: the number of send elements, it is not nbytes.
+ * @param dtype: data type
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void allGather(Comm_t comm, const void *senddata, void *recvdata, size_t sendcount, ncclDataType_t dtype,
+               cudaStream_t stream, AllGatherAlgo algo = AllGatherAlgo::kNone);
+
+/**
+ * @brief All-Reduce
+ *
+ * Reduces data arrays of length count in senddata using op operation, and
+ * leaves identical copies of result on each recvdata.
+ *
+ * In-place operation will happen if senddata == recvdata.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param count: the number of elements, it is not nbytes.
+ * @param dtype: data type
+ * @param op：Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void allReduce(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype,
+               ncclRedOp_t op, cudaStream_t stream, AllReduceAlgo algo = AllReduceAlgo::kNone);
+
+/**
+ * @brief Whether is supported non-contiguous tensors
+ * @param comm: Communicator
+ * @param dtype: data type
+ * @param shape: tensor shape
+ * @param ndim: the ndim of tensor
+ * @param numel: the number of tensor
+ * @param op: Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @return: supported
+ */
+bool allReduceStrideSupported(Comm_t comm, ncclDataType_t dtype, const int64_t *shape, int ndim, uint64_t numel, ncclRedOp_t op);
+
+/**
+ * @brief All-Reduce for non-contiguous tensors
+ * @param comm: Communicator
+ * @param senddata: send tensor
+ * @param recvdata: recv tensor
+ * @param stream: CUDA Stream
+ */
+void allReduceStride(Comm_t comm, const TensorDesc &senddata, TensorDesc &recvdata, cudaStream_t stream);
+
+/**
+ * @brief Send data from senddata to rank peer.
+ *
+ * Rank peer needs to call ncclRecv with the same datatype and the same count from this
+ * rank. This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations
+ * need to progress concurrently to complete.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param count: the number of send elements, it is not nbytes.
+ * @param dtype: data type
+ * @param peer: the destination rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void send(Comm_t comm, const void *senddata, size_t count, ncclDataType_t dtype, int peer, cudaStream_t stream,
+          SendAlgo algo = SendAlgo::kNone);
+
+/**
+ * @brief Receive data from rank peer into recvdata.
+ *
+ * Rank peer needs to call ncclSend with the same datatype and the same count to this
+ * rank. This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations
+ * need to progress concurrently to complete.
+ *
+ * @param comm: Communicator
+ * @param recvdata: recv data
+ * @param count: the number of recv elements, it is not nbytes.
+ * @param dtype: data type
+ * @param peer: source rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void recv(Comm_t comm, void *recvdata, size_t count, ncclDataType_t dtype, int peer, cudaStream_t stream,
+          RecvAlgo algo = RecvAlgo::kNone);
+
+/**
+ * @brief Reduces data arrays of length count in senddata into recvdata using op operation.
+ *
+ * Recvdata may be NULL on all calls except for root device.
+ * root is the rank (not the CUDA device) where data will reside after the
+ * operation is complete.
+ *
+ * In-place operation will happen if senddata == recvdata.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param count:  the number of elements, it is not nbytes.
+ * @param dtype: data type
+ * @param op: Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @param root: root rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void reduce(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype, ncclRedOp_t op,
+            int root, cudaStream_t stream, ReduceAlgo algo = ReduceAlgo::kNone);
+
+/**
+ * @brief Broadcast
+ *
+ * Copies count values from root to all other devices.
+ * root is the rank (not the CUDA device) where data resides before the
+ * operation is started.
+ *
+ * In-place operation will happen if senddata == recvdata.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param count: the number of elements, it is not nbytes.
+ * @param dtype: data type
+ * @param root: root rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void broadcast(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype, int root,
+               cudaStream_t stream, BroadcastAlgo algo = BroadcastAlgo::kNone);
+
+/**
+ *
+ * @brief Reduce-Scatter
+ *
+ * Reduces data in senddata using op operation and leaves reduced result
+ * scattered over the devices so that recvdata on rank i will contain the i-th
+ * block of the result.
+ * Assumes sendcount is equal to nranks*recvcount, which means that senddata
+ * should have a size of at least nranks*recvcount elements.
+ *
+ * In-place operations will happen if recvdata == senddata + rank * recvcount.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param recvcount: the number of recv elements, it is not nbytes.
+ * @param dtype: data type
+ * @param op：Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void reduceScatter(Comm_t comm, const void *senddata, void *recvdata,
+                   size_t recvcount, ncclDataType_t dtype, ncclRedOp_t op, cudaStream_t stream,
+                   ReduceScatterAlgo algo = ReduceScatterAlgo::kNone);
+
+/**
+ * @brief Send data from src_rank to dst_rank on src_rank process, recv data on dst_rank.
+ *
+ * @param comm: Communicator
+ * @param data: send data to dst rank if current rank is src_rank, recv data if current rank is dst_rank.
+ * @param count: the number of send/recv elements, it is not nbytes.
+ * @param dtype: data type
+ * @param src_rank: src rank
+ * @param dst_rank: dst rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void p2p(Comm_t comm, void *data, size_t count, ncclDataType_t dtype, int src_rank, int dst_rank, cudaStream_t stream,
+         SendAlgo algo = SendAlgo::kNone);
+
+/**
+ * @brief The all ranks of communicator send senddata to dst_rank.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdatas: recv datas，it is two-dim array, shape: [WorldSize, RecvDataPointer],
+ *        the first dim is host pointer，the second dim is GPU pointer,
+ *        it can be nullptr when current rank is not dst rank.
+ * @param sendcount: the number of send elements, it is not nbytes.
+ * @param dst_rank: dst rank
+ * @param dtype: data type
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void gather(Comm_t comm, const void *senddata, void **recvdatas, size_t sendcount, int dst_rank, ncclDataType_t dtype,
+            cudaStream_t stream, GatherAlgo algo = GatherAlgo::kNone);
+
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/comm/core/error.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <stdexcept>
+#include "status.h"
+
+namespace ixformer::comm {
+
+class CommError : public std::runtime_error {
+public:
+    template<class ERROR_STR>
+    CommError(CommStatus error, const ERROR_STR str) : error_{error}, std::runtime_error(str) {}
+
+    CommStatus status() {
+        return error_;
+    }
+
+private:
+    CommStatus error_;
+};
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/comm/core/op_algo.h

@@ -0,0 +1,80 @@
+#pragma once
+
+#include <string>
+
+namespace ixformer::comm {
+
+enum class AllGatherAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class AllReduceAlgo {
+    kNone,          // None
+    kAuto,          // 自动选择算法
+    kAllGatherSum,  // 针对小数据量的算法
+    kBroadcastSum,  // 针对小数据量的算法
+    kRing,          // Ring AllReduce
+    kQuant,         // 对通讯算法进行量化，默认使用 kQuantL1
+    kQuantL1,       // 对通讯算法进行量化，优先使用量化算法以及最大保留精度，在部分 Size 性能不佳时，退化为 Auto 算法
+    kQuantL2,       // 对通讯算法进行量化，优先使用量化算法以及最大化速度，在部分 Size 性能不佳时，退化为 Auto 算法
+    kQuantL1AllSize,// 对所有的 Size 都使用量化算法
+    kQuantL2AllSize,// 对所有的 Size 都使用量化算法
+    kNCCL,          // 使用 NCCL
+    kStride,        // 输入或输出的 Tensor 不是连续的
+    kNumAlgo
+};
+
+
+enum class BroadcastAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class GatherAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class SendAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+typedef SendAlgo RecvAlgo;
+
+enum class ReduceAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class ReduceScatterAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+
+std::string to_string(AllGatherAlgo algo);
+std::string to_string(AllReduceAlgo algo);
+std::string to_string(BroadcastAlgo algo);
+std::string to_string(GatherAlgo algo);
+std::string to_string(SendAlgo algo);
+std::string to_string(ReduceAlgo algo);
+std::string to_string(ReduceScatterAlgo algo);
+
+template<typename Algo>
+Algo get_algo_from_str(const std::string &name);
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/comm/core/status.h

@@ -0,0 +1,22 @@
+#pragma once
+#include "comm/core/common.h"
+
+
+namespace ixformer::comm {
+
+enum CommStatus {
+    commSuccess,
+    commFail,
+    commCudaError,
+    commNcclError,
+    commInvalidArgument,
+    commUnsupported,
+    commInternalError,
+    commInvalidComm// maybe comm is nullptr
+};
+
+
+std::string to_string(CommStatus status);
+
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/kernels/error.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <stdexcept>
+
+#include "status.h"
+
+namespace ixformer::kernels {
+
+class KernelError : public std::runtime_error {
+public:
+    template<class ERROR_STR>
+    KernelError(KernelStatus error, const ERROR_STR str) : error_{error}, std::runtime_error(str) {}
+
+    KernelStatus status() {
+        return error_;
+    }
+
+private:
+    KernelStatus error_;
+};
+
+}// namespace ixformer::kernels

ixformer_sdk/csrc/include/ixformer/kernels/status.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <string>
+
+namespace ixformer::kernels {
+
+enum KernelStatus {
+    kernelSuccess,
+    kernelFail,
+    kernelCudaError,
+    kernelInvalidArgument,
+    kernelCuinferError,
+    kernelUnsupported,
+};
+
+
+std::string to_string(KernelStatus status);
+
+
+}// namespace ixformer::kernels

ixformer_sdk/csrc/include/ixformer/kernels/tensor.h

@@ -0,0 +1,92 @@
+#pragma once
+
+#include <string>
+
+namespace ixformer::kernels {
+
+const uint8_t MAX_TENSOR_NDIM = 8;
+
+// align with at::ScalarType
+enum DType {
+    Byte = 0,
+    Char = 1,
+    Short = 2,
+    Int = 3,
+    Long = 4,
+    Half = 5,
+    Float = 6,
+    Double = 7,
+    ComplexHalf = 8,
+    ComplexFloat = 9,
+    ComplexDoubl = 10,
+    Bool = 11,
+    QInt8 = 12,
+    QUInt8 = 13,
+    QInt32 = 14,
+    BFloat16 = 15,
+    QUInt4x2 = 16,
+    QUInt2x4 = 17,
+    Bits1x8 = 18,
+    Bits2x4 = 19,
+    Bits4x2 = 20,
+    Bits8 = 21,
+    Bits16 = 22,
+    Float8_e5m2 = 23,
+    Float8_e4m3fn = 24,
+    Undefined = 25,
+    NumOptions = 26
+};
+
+struct TensorDesc {
+
+public:
+    // delete default constructor
+    TensorDesc() = delete;
+    // All information must be (should be) prepared when constructing a TensorDesc object.
+    TensorDesc(DType scalar_type, void *data_ptr, int64_t numel, int64_t dim, const int64_t *size, const int64_t *stride, bool is_contiguous, bool is_cuda)
+        : dtype(scalar_type), ptr(data_ptr), nnumel(numel), ndim(dim), sizes(size), strides(stride), contiguous(is_contiguous), cuda(is_cuda) {}
+
+    inline DType scalar_type() const {
+        return dtype;
+    }
+
+    inline void *data_ptr() const {
+        return ptr;
+    }
+
+    inline int64_t numel() const {
+        return nnumel;
+    }
+
+    inline int64_t dim() const {
+        return ndim;
+    }
+
+    inline int64_t size(int64_t dim) const {
+        return dim < 0 ? sizes[ndim - dim] : sizes[dim];
+    }
+
+    inline int64_t stride(int64_t dim) const {
+        return dim < 0 ? strides[ndim - dim] : strides[dim];
+    }
+
+    inline bool is_contiguous() const {
+        return contiguous;
+    }
+
+    inline bool is_cuda() const {
+        return cuda;
+    }
+
+private:
+    void *ptr{nullptr};
+    DType dtype;
+    int64_t nnumel{0};
+    int64_t ndim{0};
+    const int64_t *sizes{nullptr};
+    const int64_t *strides{nullptr};
+    bool contiguous{false};
+    bool cuda{false};
+};
+
+}// namespace ixformer::kernels

ixformer_sdk/distributed/__init__.py

@@ -0,0 +1 @@
+from ._distributed import *

ixformer_sdk/distributed/_distributed.py

@@ -0,0 +1,481 @@
+import warnings
+from collections import defaultdict
+from typing import List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+import torch.distributed.distributed_c10d as c10d
+from ixformer._C import _distributed as cdist
+from ixformer._C._distributed import comm
+from ixformer._C._distributed.comm import (
+    AllGatherAlgo,
+    AllReduceAlgo,
+    BroadcastAlgo,
+    ReduceAlgo,
+    ReduceOp,
+    ReduceScatterAlgo,
+    SendAlgo,
+)
+from ixformer.core.multi_level_cache import MultiLevelCache
+from torch import Tensor
+from torch.distributed import ProcessGroup
+
+from ixformer.core import config
+
+IxformerCommType = int
+RecvAlgo = SendAlgo
+
+_GROUP_TO_IXFC_COMM_CACHE = MultiLevelCache()
+_IXFC_COMM_TO_GROUP_CACHE = MultiLevelCache()
+
+
+def get_store(group: dist.ProcessGroup = None) -> dist.Store:
+    if group is None:
+        group = c10d._get_default_group()
+
+    return c10d._pg_map[group][1]
+
+
+class StoreWrapper(cdist.comm.C10dStoreWrapper):
+    _GROUP_COUNT = defaultdict(dict)
+
+    def __init__(self, group: ProcessGroup):
+        super().__init__()
+        self.store = get_store()
+
+        ranks = dist.get_process_group_ranks(group)
+
+        group_key = "_".join([str(r) for r in ranks])
+        if group not in self._GROUP_COUNT[group_key]:
+            self._GROUP_COUNT[group_key][group] = len(self._GROUP_COUNT[group_key])
+        group_count = self._GROUP_COUNT[group_key][group]
+
+        self.prefix = f"gid_{group_count}_" + group_key
+
+    def _gen_unique_key(self, key):
+        return f"{self.prefix}_{key}"
+
+    def set(self, key: str, value: str):
+        key = self._gen_unique_key(key)
+        self.store.set(key, value)
+
+    def get(self, key: str) -> str:
+        key = self._gen_unique_key(key)
+        self.store.wait([key])
+        return self.store.get(key).decode("utf8")
+
+
+def init_comm_with_store(group=None, shmsize: int = None):
+    if group is None:
+        group = c10d._get_default_group()
+
+    world_size = dist.get_world_size(group=group)
+    rank = dist.get_group_rank(group=group, global_rank=dist.get_rank())
+
+    if shmsize is None:
+        shmsize = config.IXFORMER_COMM_SHM_SIZE
+
+    store_wrapper = StoreWrapper(group=group)
+    ixfc_comm = cdist.comm.init_communicator_by_store(
+        store=store_wrapper, world_size=world_size, rank=rank, max_shm_mem_size=shmsize
+    )
+
+    _GROUP_TO_IXFC_COMM_CACHE.set(group, ixfc_comm)
+    _IXFC_COMM_TO_GROUP_CACHE.set(ixfc_comm, group)
+    return ixfc_comm
+
+
+_sub_store = None
+
+
+def create_nccl_unique_id(addr: str, port: str, world_size: int, rank: int):
+    global _sub_store
+    _sub_store = dist.TCPStore(
+        host_name=addr, port=int(port), world_size=world_size, is_master=rank == 0
+    )
+    store_key = "ncclUniqueId"
+    if rank == 0:
+        commid = cdist.comm.create_nccl_unique_id()
+        _sub_store.set(store_key, commid)
+    else:
+        _sub_store.wait([store_key])
+        commid = _sub_store.get(store_key).decode("utf8")
+
+    return commid
+
+
+def init_comm_with_eth(
+    addr: str, port: str, world_size: int, rank: int, shmsize: int = None
+):
+    commid = create_nccl_unique_id(addr, port, world_size=world_size, rank=rank)
+    return cdist.comm.init_communicator_by_nccl_id(commid, world_size, rank, shmsize)
+
+
+def _check_group(group: Optional[ProcessGroup] = None):
+    if group is None:
+        group = c10d._get_default_group()
+
+    if isinstance(group, ProcessGroup):
+        ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None)
+        if ixfc_comm is None:
+            return init_comm_with_store(group)
+        return ixfc_comm
+
+    return group
+
+
+def get_comm_group_stream(group: Optional[ProcessGroup] = None):
+    group = _check_group(group)
+    return comm.get_comm_group_stream(group)
+
+
+def set_comm_group_stream(stream: int, group: Optional[ProcessGroup] = None):
+    group = _check_group(group)
+    return comm.set_comm_group_stream(group, stream)
+
+
+def get_group_rank(group: Optional[ProcessGroup], global_rank) -> int:
+    """将 global rank 映射到 group 中的相对 rank"""
+    if isinstance(group, IxformerCommType):
+        _pg = _IXFC_COMM_TO_GROUP_CACHE.get(group, None)
+        if _pg is None:
+            return global_rank
+        else:
+            group = _IXFC_COMM_TO_GROUP_CACHE.get(group)
+
+    if group is None:
+        group = c10d._get_default_group()
+
+    return dist.get_group_rank(group, global_rank)
+
+
+def get_global_rank(group: Optional[ProcessGroup], group_rank: int) -> int:
+    """将一个 group rank 映射到 global rank"""
+    if group is None:
+        group = c10d._get_default_group()
+    return c10d.get_global_rank(group, group_rank)
+
+
+def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> List[int]:
+    """获取 Group 的 global ranks"""
+    if group is None:
+        group = c10d._get_default_group()
+    return c10d.get_process_group_ranks(group)
+
+
+def new_group(ranks: List[int] = None, shmsize=None, *args, **kwargs):
+    """通过 global ranks 去创建一个通讯组"""
+    group = c10d.new_group(ranks, *args, **kwargs)
+
+    if ranks is None:
+        ranks = dist.get_process_group_ranks(group)
+
+    if get_rank() in ranks:
+        init_comm_with_store(group=group, shmsize=shmsize)
+    return group
+
+
+def new_subgroups_by_enumeration(
+    ranks_per_subgroup_list, shmsize=None, *args, **kwargs
+) -> Tuple[ProcessGroup, List[ProcessGroup]]:
+    """
+    通过一组 global ranks 去创建通讯组
+
+    :param ranks_per_subgroup_list: global ranks
+    :return: 返回当前 rank 所在的通讯组 和 新的 subgroups
+    """
+    self_group, other_group = c10d.new_subgroups_by_enumeration(
+        ranks_per_subgroup_list, *args, **kwargs
+    )
+    init_comm_with_store(self_group, shmsize=shmsize)
+    return self_group, other_group
+
+
+def destroy_process_group(group: Optional[ProcessGroup] = None):
+    """销毁 Group"""
+    if group is None:
+        group = c10d._get_default_group()
+    ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None)
+
+    if ixfc_comm is None:
+        dist.destroy_process_group(group)
+    else:
+        comm.destroy(ixfc_comm)
+        dist.destroy_process_group(group)
+
+
+def get_rank(group: Optional[ProcessGroup] = None) -> int:
+    """获取当前进程的 Rank，如果 group 是 null，那么返回的是 Global Rank, 否则返回的相对的 Rank，即在当前组中的 rank"""
+    return c10d.get_rank(group)
+
+
+def get_world_size(group: Optional[ProcessGroup] = None) -> int:
+    """获取 Group 中的成员大小"""
+    return c10d.get_world_size(group)
+
+
+def barrier(group: Optional[ProcessGroup] = None, use_comm_stream: bool = False):
+    """同步 Group 中的 rank"""
+    group = _check_group(group)
+    comm.barrier(group, use_comm_stream)
+
+
+def isend(
+    tensor: Tensor,
+    dst: int,
+    group: Optional[ProcessGroup] = None,
+    use_comm_stream: bool = False,
+):
+    dst = get_group_rank(group, dst)
+    group = _check_group(group)
+    return comm.send(group, tensor, dst, use_comm_stream, SendAlgo.kNone)
+
+
+def send(*args, **kwargs):
+    warnings.warn("not support sync mode, as async to call.")
+    return isend(*args, **kwargs)
+
+
+def irecv(
+    tensor: torch.Tensor,
+    src: int,
+    group: Optional[ProcessGroup] = None,
+    use_comm_stream: bool = False,
+):
+    src = get_group_rank(group, src)
+    group = _check_group(group)
+    return comm.recv(group, tensor, src, use_comm_stream, SendAlgo.kNone)
+
+
+def recv(*args, **kwargs):
+    warnings.warn("not support sync mode, as async to call.")
+    return irecv(*args, **kwargs)
+
+
+def point_to_point(
+    tensor: Tensor,
+    src: int,
+    dst: int,
+    group: Optional[ProcessGroup] = None,
+    use_comm_stream: bool = False,
+):
+    """在 src rank 发送 tensor，在 dst_rank 上接收数据到 tensor 中"""
+    src = get_group_rank(group, src)
+    dst = get_group_rank(group, dst)
+    group = _check_group(group)
+    return comm.p2p(group, tensor, src, dst, use_comm_stream)
+
+
+def reduce(
+    tensor,
+    root: int,
+    op=ReduceOp.SUM,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    out: Tensor = None,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        ixf_tensor = torch.tensor([1], device="cuda")
+        ixfd.reduce(ixf_tensor, 1, async_op=True)
+        print("rank {rank}:", ixf_tensor)
+
+        # output
+        rank 0:  tensor([1], device='cuda:0')
+        rank 1:  tensor([4], device='cuda:1')
+        rank 2:  tensor([1], device='cuda:2')
+        rank 3:  tensor([1], device='cuda:3')
+    """
+
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    if out is None:
+        out = tensor
+
+    root = get_group_rank(group, root)
+    group = _check_group(group)
+    return comm.reduce(group, tensor, out, op, root, use_comm_stream, ReduceAlgo.kNone)
+
+
+def broadcast(
+    tensor: Tensor,
+    src: int,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    out: Tensor = None,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        ixf_tensor = torch.tensor([rank], device="cuda")
+        ixfd.broadcast(ixf_tensor, 1, async_op=True)
+        print("rank {rank}: ", ixf_tensor)
+
+        # output
+        rank 0:  tensor([1], device='cuda:0')
+        rank 1:  tensor([1], device='cuda:1')
+        rank 2:  tensor([1], device='cuda:2')
+        rank 3:  tensor([1], device='cuda:3')
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    if out is None:
+        out = tensor
+
+    src = get_group_rank(group, src)
+    group = _check_group(group)
+    return comm.broadcast(group, tensor, out, src, use_comm_stream, BroadcastAlgo.kNone)
+
+
+def reduce_scatter_tensor(
+    output: Tensor,
+    input: Tensor,
+    op=ReduceOp.SUM,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        ixf_tensor_out = torch.zeros(2, dtype=torch.int64, device="cuda")
+        tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device="cuda")
+        # tensor_in: tensor([0, 1, 2, 3, 4, 5, 6, 7], device='cuda:0')
+
+        ixfd.reduce_scatter_tensor(ixf_tensor_out, tensor_in, async_op=True)
+        print("rank {rank}:", ixf_tensor_out)
+
+        # output
+        rank 0:  tensor([0, 4], device='cuda:0')
+        rank 1:  tensor([ 8, 12], device='cuda:1')
+        rank 2:  tensor([16, 20], device='cuda:2')
+        rank 3:  tensor([24, 28], device='cuda:3')
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    group = _check_group(group)
+    return comm.reduce_scatter(
+        group, input, output, op, use_comm_stream, ReduceScatterAlgo.kNone
+    )
+
+
+def all_reduce(
+    tensor: Tensor,
+    op=ReduceOp.SUM,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    out: Tensor = None,
+    algo: AllReduceAlgo = AllReduceAlgo.kNone,
+    use_comm_stream: bool = False,
+):
+    """
+    Args:
+        tensor: inpute tensor
+        op: ReduceOp: SUM, MIN or MAX
+        group: communicator group
+        async_op: ixformer support async mode
+        out: output tensor
+        algo: AllReduce Algo: Auto, Quant, QuantL1, QuantL2, NCCL, Ring, AllGatherSum, BroadcastSum
+        use_comm_stream: ixformer support set communication stream by ixformer.distributed.set_comm_group_stream,
+          if true, submit the kernels of communication to communication stream,
+          if false, use current stream by torch.cuda.current_stream
+    Returns: out
+
+    Example:
+        >>> # All tensors below are of torch.int64 type.
+        >>> # We have 2 process groups, 2 ranks.
+        >>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
+        >>> tensor
+        tensor([1, 2]) # Rank 0
+        tensor([3, 4]) # Rank 1
+        >>> ixfd.all_reduce(tensor, op=ReduceOp.SUM, async_op=True)
+        >>> tensor
+        tensor([4, 6]) # Rank 0
+        tensor([4, 6]) # Rank 1
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    group = _check_group(group)
+
+    if out is None:
+        out = tensor
+
+    comm.all_reduce(
+        group,
+        tensor,
+        out,
+        op,
+        use_comm_stream=use_comm_stream,
+        algo=algo,
+    )
+
+
+def all_gather_into_tensor(
+    output: Tensor,
+    input: Tensor,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        tensor_in = torch.arange(2, dtype=torch.int64, device="cuda") + 1 + 2 * rank
+            rank 0:  tensor in:  tensor([1, 2], device='cuda:0')
+            rank 1:  tensor in:  tensor([3, 4], device='cuda:1')
+            rank 2:  tensor in:  tensor([5, 6], device='cuda:2')
+            rank 3:  tensor in:  tensor([7, 8], device='cuda:3')
+
+        ixf_tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device="cuda")
+        ixfd.all_gather_into_tensor(ixf_tensor_out, tensor_in, async_op=True)
+        print("rank {rank}:", ixf_tensor_out)
+
+        # output:
+        rank 0:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0')
+        rank 1:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:1')
+        rank 2:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:2')
+        rank 3:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:3')
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    group = _check_group(group)
+    return comm.all_gather(
+        group, input, output, use_comm_stream, algo=AllGatherAlgo.kNone
+    )
+
+
+def gather(
+    tensor,
+    gather_list=None,
+    dst=0,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        >>> # We have 2 process groups, 2 ranks.
+        >>> tensor = torch.tensor(rank+1,dtype=torch.float32).cuda()
+        >>> tensor
+        tensor(1.) # Rank 0
+        tensor(2.) # Rank 1
+        >>> gather_list = [torch.zeros(1).cuda() for _ in range(rank)] if rank == dst else None
+        >>> gather_list
+        [tensor([0,]),tensor([1,])] # Rank 0
+        None                        # Rank 1
+        ixfd.gather(tensor,gather_list,0,async_op=True)
+        >>> gather_list
+        [tensor([1.]),tensor([2.])] # Rank 0
+        None                        # Rank 1
+    """
+    gather_list = gather_list if gather_list is not None else []
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    dst = get_group_rank(group, dst)
+    group = _check_group(group)
+    return comm.gather(group, tensor, gather_list, dst, use_comm_stream)

ixformer_sdk/distributed/overlap_comm.py

@@ -0,0 +1,412 @@
+import abc
+import enum
+import os
+from contextlib import contextmanager, nullcontext
+from typing import List, Optional
+
+import torch.cuda
+from ixformer.core.dispatcher import Dispatcher
+
+from ixformer.core import config
+
+from . import _distributed as ixfd
+
+
+class SplitOverlapComm(Dispatcher):
+    def __init__(self, num_chunks, num_compute_streams=None, comm_group=None):
+        """
+        Args:
+            num_chunks: the number of chunks
+            num_compute_streams: the number of compute streams, default: 1
+            comm_group: communicator group
+        """
+
+        self._num_chunks = num_chunks
+        self._num_compute_streams = num_compute_streams or 1
+        self._comm_group = comm_group
+
+        self._compute_streams: List[torch.cuda.Stream] = self.create_compute_streams()
+        self._comm_stream: torch.cuda.Stream = torch.cuda.Stream(priority=-1)
+
+        self._start_compute_event: torch.cuda.Event = torch.cuda.Event()
+        self._stop_compute_event: torch.cuda.Event = torch.cuda.Event()
+
+        self._start_comm_event: torch.cuda.Event = torch.cuda.Event()
+        self._stop_comm_event: torch.cuda.Event = torch.cuda.Event()
+
+        # keep origin state
+        self._main_stream: Optional[torch.cuda.Stream] = None
+        self._origin_ixf_comm_stream = None
+        self._ixformer_streams = dict()
+
+    @classmethod
+    def dispatcher_key(
+        cls, num_chunks, num_compute_streams=None, comm_group=None, *args, **kwargs
+    ):
+        """
+        the key of SplitOverlapComm
+        Args:
+            num_chunks: the number of chunks
+            num_compute_streams: the number of compute streams, default: 1
+            comm_group: communicator group
+        Returns: unique key
+        """
+        # warn: keey same function parameters with init
+        return (cls.__name__, num_chunks, num_compute_streams, comm_group)
+
+    @classmethod
+    def enable(cls):
+        return config.IXFORMER_ENABLE_OVERLAP_COMM
+
+    @property
+    def num_chunks(self):
+        return self._num_chunks
+
+    @property
+    def num_compute_streams(self):
+        return self._num_compute_streams
+
+    @property
+    def comm_group(self):
+        return self._comm_group
+
+    def create_compute_streams(self):
+        streams = []
+        for _ in range(self.num_compute_streams):
+            streams.append(torch.cuda.Stream())
+        return streams
+
+    def start_overlap(self):
+        self._main_stream = torch.cuda.current_stream()
+
+        self._start_compute_event.record(torch.cuda.current_stream())
+        for compute_stream in self._compute_streams:
+            compute_stream.wait_event(self._start_compute_event)
+
+        self._origin_ixf_comm_stream = ixfd.get_comm_group_stream(self._comm_group)
+        ixfd.set_comm_group_stream(self._comm_stream.cuda_stream, self._comm_group)
+
+    def stop_overlap(self):
+        last_compute_stream_id = (
+            self.num_chunks + self.num_compute_streams - 1
+        ) % self.num_compute_streams
+        self._stop_compute_event.record(self._compute_streams[last_compute_stream_id])
+        self._stop_comm_event.record(self._comm_stream)
+        torch.cuda.current_stream().wait_event(self._stop_compute_event)
+        torch.cuda.current_stream().wait_event(self._stop_comm_event)
+
+        ixfd.set_comm_group_stream(self._origin_ixf_comm_stream, self._comm_group)
+
+    def start_comm(self, chunk_idx):
+        """
+        prepare communication stream and wait event.
+        Args:
+            chunk_idx: the index of chunk
+        """
+
+        self._start_comm_event.record(
+            self._compute_streams[chunk_idx % self.num_compute_streams]
+        )
+        self._comm_stream.wait_event(self._start_comm_event)
+
+    @contextmanager
+    def compute_stream_context(self, chunk_idx):
+        """
+        open python context and switch to compute stream in torch context
+        Args:
+            chunk_idx: the index of chunk
+        """
+
+        stream = self._compute_streams[chunk_idx % self.num_compute_streams]
+
+        # print("before stream:", torch.cuda.current_stream())
+        torch.cuda.set_stream(stream)
+
+        # print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream())
+        yield stream
+
+        torch.cuda.set_stream(self._main_stream)
+
+    @contextmanager
+    def stream_context(self, stream):
+        # print("before stream:", torch.cuda.current_stream())
+        torch.cuda.set_stream(stream)
+
+        # print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream())
+        yield stream
+
+        torch.cuda.set_stream(self._main_stream)
+
+    def forward(self, *args, **kwargs):
+        self.start_overlap()
+        out = self.compute(*args, **kwargs)
+        self.stop_overlap()
+        return out
+
+    @abc.abstractmethod
+    def compute(self, *args, **kwargs):
+        """
+        it is abstract method to execute compute and communication.
+        """
+        pass
+
+
+class GemmMethod(enum.IntEnum):
+    kCUINFER = 0
+    kCUBLAS = 1
+    kLIMITED_GEMM = 2
+
+
+class GemmWithLimitedBlock:
+    def __init__(self, limit_algo=0) -> None:
+        self.limit_algo = limit_algo
+        self.env_key = "PYTORCH_GEMM_BLOCK_LIMITATION"
+
+    def __enter__(self) -> None:
+        os.environ[self.env_key] = str(self.limit_algo)
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        del os.environ[self.env_key]
+
+
+class IxFormerLimitedGemmContext:
+    def __init__(self) -> None:
+        self.env_key = "IXFORMER_ENABLE_PERSISTENT_GEMM"
+
+    def __enter__(self) -> None:
+        os.environ[self.env_key] = "1"
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        os.environ[self.env_key] = "0"
+
+
+class GemmAllReduceSplitOverlapComm(SplitOverlapComm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.gemm_method_env = config.IXFORMER_OVERLAP_GEMM_METHOD
+
+        if self.gemm_method_env is None:
+            if ixfd.get_world_size(self.comm_group) == 2:
+                self.gemm_method_env = 0
+            else:
+                self.gemm_method_env = 2
+
+        self.gemm_method = GemmMethod(int(self.gemm_method_env))
+        self.limited_gemm_ctx = GemmWithLimitedBlock()
+        self.ixf_limited_gemm_ctx = IxFormerLimitedGemmContext()
+        self.split_ratio = config.IXFORMER_OVERLAP_SPLIT_RATIO
+
+    @classmethod
+    def compute_row_parallel_dims(cls, input):
+        batch = 1
+        if input.ndim == 2:
+            seqlen = input.shape[0]
+        else:
+            batch = input.shape[0]
+            seqlen = input.shape[1]
+
+        parallel_dims = batch * seqlen
+        return parallel_dims
+
+    def compute(self, input, weight, bias=None, out=None, *args, **kwargs):
+        """
+        :param input: [Batch, SeqLen, Hidden]
+        :param weight: [OutChannel, InChannel]
+        :param bias: [OutChannel]
+        """
+
+        is_update_shape = input.ndim > 2
+        batch = 1
+        if input.ndim == 2:
+            seqlen = input.shape[0]
+        else:
+            batch = input.shape[0]
+            seqlen = input.shape[1]
+
+        parallel_dims = batch * seqlen
+
+        if is_update_shape:
+            input = input.reshape(parallel_dims, -1)
+
+        if out is None:
+            out_shape = [parallel_dims, weight.shape[0]]
+            out_dtype = kwargs["out_dtype"] if "out_dtype" in kwargs else input.dtype
+            out = torch.empty(out_shape, dtype=out_dtype, device=input.device)
+
+        if self.split_ratio is not None:
+            round_multiples = 256 if parallel_dims >= 256 else parallel_dims
+            first_chunk_size = (
+                round((parallel_dims * float(self.split_ratio)) / round_multiples)
+                * round_multiples
+            )
+            middle_chunk_size = (parallel_dims - first_chunk_size) // (
+                self.num_chunks - 1
+            )
+            middle_chunk_size = (middle_chunk_size // round_multiples) * round_multiples
+            last_chunk_size = (
+                parallel_dims
+                - first_chunk_size
+                - middle_chunk_size * (self.num_chunks - 2)
+            )
+
+            chunk_sizes = (
+                [first_chunk_size]
+                + [middle_chunk_size] * (self.num_chunks - 2)
+                + [last_chunk_size]
+            )
+            input_chunks = torch.split_with_sizes(input, chunk_sizes, dim=0)
+            out_chunks = torch.split_with_sizes(out, chunk_sizes, dim=0)
+
+            # print(first_chunk_size, middle_chunk_size, last_chunk_size, chunk_sizes)
+        else:
+            input_chunks = torch.chunk(input, self.num_chunks, dim=0)
+            out_chunks = torch.chunk(out, self.num_chunks, dim=0)
+
+        for chunk_idx in range(len(input_chunks)):
+            with self.compute_stream_context(chunk_idx):
+                chunk_out = self.gemm_dispatcher(
+                    chunk_idx,
+                    input_chunks[chunk_idx],
+                    weight,
+                    out_chunks[chunk_idx],
+                    *args,
+                    **kwargs,
+                )
+
+            self.start_comm(chunk_idx)
+
+            ixfd.all_reduce(
+                chunk_out, async_op=True, group=self.comm_group, use_comm_stream=True
+            )
+
+        if is_update_shape:
+            out = out.reshape(batch, seqlen, -1)
+
+        if bias is not None:
+            out = out + bias
+
+        return out
+
+    def gemm_dispatcher(
+        self,
+        chunk_idx,
+        chunk_input,
+        weight,
+        chunk_out=None,
+        user_gemm_method=None,
+        *args,
+        **kwargs,
+    ):
+        if user_gemm_method is not None and callable(user_gemm_method):
+            ctx = nullcontext() if chunk_idx == 0 else self.ixf_limited_gemm_ctx
+            with ctx:
+                return user_gemm_method(
+                    chunk_input, weight, out=chunk_out, *args, **kwargs
+                )
+
+        if user_gemm_method is None:
+            user_gemm_method = self.gemm_method
+
+        if user_gemm_method == GemmMethod.kCUINFER:
+            import ixformer.functions as ixff
+
+            return ixff.linear(chunk_input, weight, output=chunk_out)
+        elif user_gemm_method == GemmMethod.kCUBLAS:
+            return torch.matmul(chunk_input, weight.T, out=chunk_out)
+        elif user_gemm_method == GemmMethod.kLIMITED_GEMM:
+            ctx = self.limited_gemm_ctx
+            with ctx:
+                return torch.matmul(chunk_input, weight.T, out=chunk_out)
+        elif user_gemm_method == GemmMethod.kCUBLAS:
+            return torch.matmul(chunk_input, weight.T, out=chunk_out)
+        else:
+            raise RuntimeError(f"Invalid gemm method, got {self.gemm_method}.")
+
+    @classmethod
+    def native_forward(
+        cls,
+        input,
+        weight,
+        bias=None,
+        out=None,
+        group=None,
+        user_gemm_method=None,
+        *args,
+        **kwargs,
+    ):
+        if user_gemm_method is not None and callable(user_gemm_method):
+            gemm_out = user_gemm_method(
+                input, weight, bias=bias, out=out, *args, **kwargs
+            )
+            out = out if gemm_out is None else gemm_out
+        else:
+            import ixformer.functions as ixff
+
+            # warning: 下面的两种 gemm 可能存在精度不一致
+            # out = torch.matmul(input, weight.T, out=out)
+            out = ixff.linear(input=input, weight=weight, bias=bias, output=out)
+        ixfd.all_reduce(out, async_op=True, group=group)
+        return out
+
+    @classmethod
+    def is_supported(cls, input, num_chunks, comm_group):
+        if not cls.enable():
+            return False
+
+        ndim = input.ndim
+        shape = input.shape
+
+        if ndim == 1:
+            m, k = 1, shape[0]
+        elif ndim == 2:
+            m, k = shape
+        else:
+            m, k = sum(shape[:-1]), shape[-1]
+
+        return m >= 512
+
+
+_DEFAULT_OVERLAP_GROUP = None
+_DEFAULT_OVERLAP_COMM_N2 = None
+_DEFAULT_OVERLAP_COMM_N4 = None
+_DEFAULT_OVERLAP_CHUNKS = config.IXFORMER_OVERLAP_CHUNKS
+
+
+def linear_allreduce_overlap(
+    input, weight, bias=None, out=None, group=None, num_chunks=None, *args, **kwargs
+):
+    num_chunks = num_chunks or _DEFAULT_OVERLAP_CHUNKS
+
+    # print("call overlap:", GemmAllReduceSplitOverlapComm.is_supported(input, num_chunks=num_chunks, comm_group=group), input.shape, weight.shape if torch.is_tensor(weight) else None, "WorldSize:", ixfd.get_group_world_size(group), ", NumChunks:", num_chunks)
+    if not GemmAllReduceSplitOverlapComm.is_supported(
+        input, num_chunks=num_chunks, comm_group=group
+    ):
+        return GemmAllReduceSplitOverlapComm.native_forward(
+            input, weight, bias=bias, out=out, group=group, *args, **kwargs
+        )
+
+    global _DEFAULT_OVERLAP_GROUP
+    global _DEFAULT_OVERLAP_COMM_N2
+    global _DEFAULT_OVERLAP_COMM_N4
+
+    if _DEFAULT_OVERLAP_GROUP is None:
+        _DEFAULT_OVERLAP_GROUP = group
+
+    if num_chunks == 2 and group == _DEFAULT_OVERLAP_GROUP:
+        if _DEFAULT_OVERLAP_COMM_N2 is None:
+            _DEFAULT_OVERLAP_COMM_N2 = GemmAllReduceSplitOverlapComm.dispatcher(
+                num_chunks=num_chunks, comm_group=group
+            )
+        overlap_comm = _DEFAULT_OVERLAP_COMM_N2
+    elif num_chunks == 4 and group == _DEFAULT_OVERLAP_GROUP:
+        if _DEFAULT_OVERLAP_COMM_N4 is None:
+            _DEFAULT_OVERLAP_COMM_N4 = GemmAllReduceSplitOverlapComm.dispatcher(
+                num_chunks=num_chunks, comm_group=group
+            )
+        overlap_comm = _DEFAULT_OVERLAP_COMM_N4
+    else:
+        overlap_comm = GemmAllReduceSplitOverlapComm.dispatcher(
+            num_chunks=num_chunks, comm_group=group
+        )
+    return overlap_comm.forward(input, weight, bias=bias, out=out, *args, **kwargs)

ixformer_sdk/functions/__init__.py

@@ -0,0 +1 @@
+from ..inference.functions import *

ixformer_sdk/inference/distributed/__init__.py

@@ -0,0 +1 @@
+from .mpi_utils import *

ixformer_sdk/inference/distributed/mpi_utils.py

@@ -0,0 +1,21 @@
+import os
+
+from mpi4py import MPI
+
+
+def get_world_size(comm=None):
+    if comm is None:
+        comm = MPI.COMM_WORLD
+
+    return comm.Get_size()
+
+
+def get_local_rank(comm=None):
+    return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
+
+
+def get_rank(comm=None):
+    if comm is None:
+        comm = MPI.COMM_WORLD
+
+    return comm.Get_rank()

ixformer_sdk/inference/functions/__init__.py

@@ -0,0 +1,44 @@
+from .act_and_mul import *
+from .act_bias_mm import *
+from .add import *
+from .bert import *
+from .bnb_dequant import *
+from .bnb_double_quant import *
+from .bnb_mm_dequant import *
+from .bnb_qgemm import *
+from .bnb_quant import *
+from .bnb_rowcol_absmax import *
+from .conv2d import *
+from .cross_entropy_loss import *
+from .flash_attn import *
+from .flash_attn_lib import *
+from .fused_rope import *
+from .gemv import *
+from .groupnorm import *
+from .i8w8o32 import *
+from .layernorm import *
+from .lightllm import *
+from .linalg import *
+from .linear import *
+from .lmdeploy import *
+from .marlin import *
+from .matmul import *
+from .mla_fused import *
+from .mm import *
+from .moe import *
+from .overlap_comm import *
+from .paged_attention import *
+from .quantized_linear import *
+from .residual_bias import *
+from .rms_norm import *
+from .scaled_dot_product_attention import *
+from .smoothquant import *
+from .softmax import *
+from .store_kv_cache import *
+from .t5 import *
+from .tgi import *
+from .vllm import *
+from .w8a8 import *
+from .w8a16 import *
+from .wi4a16 import *
+from .wui4a16 import *

ixformer_sdk/inference/functions/act_and_mul.py

@@ -0,0 +1,88 @@
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+import torch.nn.functional as NNF
+
+__all__ = ["ref_silu_and_mul", "ref_gelu_and_mul", "ref_gelu_tanh_and_mul",
+           "silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul"]
+
+
+def ref_silu_and_mul(input: "torch.Tensor") -> torch.Tensor:
+    x1, x2 = input.chunk(chunks=2, dim=-1)
+    res = NNF.silu(x1) * x2
+    return res
+
+
+def ref_gelu_and_mul(input: "torch.Tensor", gate_first=True) -> torch.Tensor:
+    x1, x2 = input.chunk(chunks=2, dim=-1)
+    if gate_first:
+        res = NNF.gelu(x1) * x2
+    else:
+        res = NNF.gelu(x2) * x1
+    return res
+
+
+def ref_gelu_tanh_and_mul(input: "torch.Tensor") -> torch.Tensor:
+    x1, x2 = input.chunk(chunks=2, dim=-1)
+    res = NNF.gelu(x1) * x2
+    return res
+
+
+def silu_and_mul(input: torch.Tensor, output: torch.Tensor = None):
+    
+    """
+    Args:
+        input:              (..., 2*hidden_size)    torch.float16, torch.bfloat16, torch.float32
+        output:             (..., hidden_size)      torch.float16, torch.bfloat16, torch.float32
+        
+    Returns:
+        output:             (..., hidden_size)      torch.float16, torch.bfloat16, torch.float32
+    """
+    if output is None:
+        output_shape = list(input.shape)
+        output_shape[-1] = output_shape[-1] // 2
+        output = input.new_empty(output_shape)
+    
+    ops.infer.silu_and_mul(input, output)
+    
+    return output
+
+
+def gelu_and_mul(input: "torch.Tensor", output: torch.Tensor = None, gate_first=True):
+    
+    """
+    Args:
+        input:              (..., 2*hidden_size)    torch.float16, torch.bfloat16, torch.float32
+        output:             (..., hidden_size)      torch.float16, torch.bfloat16, torch.float32
+        gate_first:                                 bool
+    Returns:
+        output:             (..., hidden_size)      torch.float16, torch.bfloat16, torch.float32
+    """
+    if output is None:
+        output_shape = list(input.shape)
+        output_shape[-1] = output_shape[-1] // 2
+        output = input.new_empty(output_shape)
+
+    ops.infer.gelu_and_mul(input, output, gate_first)
+
+    return output
+
+
+def gelu_tanh_and_mul(input: torch.Tensor, output: torch.Tensor = None):
+    
+    """
+    Args:
+        input:              (..., 2*hidden_size)    torch.float16, torch.bfloat16, torch.float32
+        output:             (..., hidden_size)      torch.float16, torch.bfloat16, torch.float32
+    Returns:
+        output:             (..., hidden_size)      torch.float16, torch.bfloat16, torch.float32
+    """
+    if output is None:
+        output_shape = list(input.shape)
+        output_shape[-1] = output_shape[-1] // 2
+        output = input.new_empty(output_shape)
+
+    ops.infer.gelu_tanh_and_mul(input, output)
+
+    return output

ixformer_sdk/inference/functions/act_bias_mm.py

@@ -0,0 +1,89 @@
+import ixformer._C as ops
+import torch
+import torch.nn.functional as NNF
+
+__all__ = ["act_bias_mm", "ref_act_bias_mm"]
+
+
+def ref_act_bias_mm(
+    mat1: torch.Tensor,
+    mat2: torch.Tensor,
+    bias: torch.Tensor = None,
+    scale: float = 1,
+    act_type: str = "none",
+    trans_format: str = "NN",
+):
+    assert len(mat1.shape) >= 2
+    assert len(mat2.shape) >= 2
+    if trans_format == "NN":
+        if bias is not None:
+            output = torch.matmul(mat1, mat2) * scale + bias
+        else:
+            output = torch.matmul(mat1, mat2) * scale
+    else:
+        if bias is not None:
+            output = torch.matmul(mat1, mat2.transpose(-1, -2)) * scale + bias
+        else:
+            output = torch.matmul(mat1, mat2.transpose(-1, -2)) * scale
+    if act_type == "gelu":
+        output = NNF.gelu(output)
+    elif act_type == "relu":
+        output = NNF.relu(output)
+    elif act_type == "silu":
+        output = NNF.silu(output)
+    elif act_type == "none":
+        output = output
+    else:
+        raise NotImplementedError()
+    return output
+
+
+def act_bias_mm(
+    mat1: torch.Tensor,
+    mat2: torch.Tensor,
+    bias: torch.Tensor = None,
+    output: torch.Tensor = None,
+    scale: float = 1,
+    act_type: str = "none",
+    trans_format: str = "NN",
+):
+    """
+    Args:
+        mat1:               [m,k] or [batch_count,m,k]          torch.float16
+        mat2:               [k,n] or [n,k]                      torch.float16
+                当trans_format为"NN"时[k,n], 当trans_format为"TN"时[n,k]
+        bias:               [n]                                 torch.float16
+        output:             [m,n]                               torch.float16
+        scale:                                                  float
+        act_type:           silu/gelu/relu/None                 str
+                如果act_type不为None,则bias也不可以为None
+        trans_format:       NN or TN                            str
+    Returns:
+        output:             [m,n]                               torch.float16
+    """
+    assert len(mat1.shape) >= 2
+    assert len(mat2.shape) >= 2
+    if output is None:
+        output_shape = list(mat1.shape)
+        m = mat1.shape[-2]
+        if trans_format == "NN":
+            n = mat2.shape[-1]
+        else:
+            n = mat2.shape[-2]
+        output_shape[-2] = m
+        output_shape[-1] = n
+        output = mat1.new_empty(output_shape)
+
+    add_bias = False
+    if bias is not None:
+        add_bias = True
+
+    if add_bias:
+        ops.infer.act_bias_mm(
+            mat1, mat2, bias, output, add_bias, scale, act_type, trans_format
+        )
+    else:
+        ops.infer.act_bias_mm(
+            mat1, mat2, mat1, output, add_bias, scale, act_type, trans_format
+        )
+    return output

ixformer_sdk/inference/functions/add.py

@@ -0,0 +1,46 @@
+import ixformer._C as ops
+import torch
+
+__all__ = [
+    "ref_add",
+    "add",
+]
+
+
+def ref_add(input: torch.Tensor, other: torch.Tensor, out: torch.Tensor = None):
+    return torch.add(input, other, out=out)
+
+
+def add(input: torch.Tensor, other: torch.Tensor, out: torch.Tensor = None):
+    """
+    out = input + other
+    Support elementwise addition, but broadcasting is not supported yet.
+    Note: The dtype of input and other needs to be the same.
+    Args:
+        input:              (...)        torch.float32, torch.float16, torch.bfloat16
+        other:              (...)        same as input
+        out:                (...)        same as input
+    Returns:
+        out:                (...)        same as input
+    """
+    if input.dtype not in [torch.float16, torch.float32, torch.bfloat16]:
+        return torch.add(input, other, out=out)
+    if not input.is_contiguous() or not other.is_contiguous():
+        return torch.add(input, other, out=out)
+    if out is not None and not out.is_contiguous():
+        return torch.add(input, other, out=out)
+
+    if input.dtype != other.dtype:
+        return torch.add(input, other, out=out)
+    if out is not None and out.dtype != input.dtype:
+        return torch.add(input, other, out=out)
+
+    assert input.shape == other.shape, (f"broadcasting is not supported yet."
+        "input is {input.shape}, other is {other.shape}")
+    
+    if out is None:
+        out = torch.empty_like(input)
+
+    ops.infer.add(input, other, out)
+
+    return out

ixformer_sdk/inference/functions/bert.py

@@ -0,0 +1,199 @@
+import ixformer._C as ops
+import torch
+
+__all__ = [
+    "ref_bert_embedding",
+    "bert_embedding",
+    "ref_bert_add_norm",
+    "bert_add_norm",
+    "ref_bert_unpack_start_end_logits",
+    "bert_unpack_start_end_logits",
+    "ref_bert_linear_residual",
+    "bert_linear_residual",
+]
+
+
+def ref_bert_embedding(
+    token_weight: torch.Tensor,
+    pos_weight: torch.Tensor,
+    type_weight: torch.Tensor,
+    ln_weight: torch.Tensor,
+    ln_bias: torch.Tensor,
+    token_ids: torch.Tensor,
+    pos_ids: torch.Tensor,
+    type_ids: torch.Tensor,
+    epsilon: float = 1e-5,
+    out: torch.Tensor = None,
+):
+    assert out is None
+    emd1 = torch.nn.functional.embedding(token_ids, token_weight)
+    emd2 = torch.nn.functional.embedding(pos_ids, pos_weight)
+    emd3 = torch.nn.functional.embedding(type_ids, type_weight)
+
+    out = emd1 + emd2 + emd3
+    out = torch.nn.functional.layer_norm(out, [out.shape[-1]], ln_weight, ln_bias)
+    return out
+
+
+def bert_embedding(
+    token_weight: torch.Tensor,
+    pos_weight: torch.Tensor,
+    type_weight: torch.Tensor,
+    ln_weight: torch.Tensor,
+    ln_bias: torch.Tensor,
+    token_ids: torch.Tensor,
+    pos_ids: torch.Tensor,
+    type_ids: torch.Tensor,
+    epsilon: float = 1e-5,
+    out: torch.Tensor = None,
+):
+    """
+    Args:
+        token_weight:       (vocab_size, hidden_size)    torch.float16, torch.bfloat16
+        pos_weight:         (pos_size, hidden_size)      same as token_weight
+        type_weight:        (type_size, hidden_size)     same as token_weight
+        ln_weight:          (hidden_size)                same as token_weight
+        ln_bias:            (hidden_size)                same as token_weight
+        token_ids:          (num_tokens)                 torch.int32, torch.int64
+        pos_ids:            (num_tokens)                 same as token_ids
+        type_ids:           (num_tokens)                 same as token_ids
+        epsilon:                                         float
+        out:                (num_tokens, hidden_size)    same as token_weight
+    Returns:
+        out:                (num_tokens, hidden_size)    same as token_weight
+    """
+    if out is None:
+        out_shape = list(token_ids.shape)
+        hidden_size = token_weight.shape[-1]
+        out_shape.append(hidden_size)
+        out = token_weight.new_empty(out_shape)
+
+    ops.infer.bert_embedding(
+        token_weight,
+        pos_weight,
+        type_weight,
+        ln_weight,
+        ln_bias,
+        token_ids,
+        pos_ids,
+        type_ids,
+        out,
+        epsilon,
+    )
+
+    return out
+
+
+def ref_bert_add_norm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    ln_weight: torch.Tensor,
+    ln_bias: torch.Tensor,
+    epsilon: float = 1e-5,
+    out: torch.Tensor = None,
+):
+    assert out is None
+    input = input + residual
+    return torch.nn.functional.layer_norm(
+        input, [input.shape[-1]], ln_weight, ln_bias, epsilon
+    )
+
+
+def bert_add_norm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    ln_weight: torch.Tensor,
+    ln_bias: torch.Tensor,
+    epsilon: float = 1e-5,
+    out: torch.Tensor = None,
+):
+    """
+    out = input + residual
+    out = add_norm(out, ln_weight, ln_bias, epsilon)
+    Args:
+        input:       (num_tokens, hidden_size)    torch.float16, torch.bfloat16
+        residual:    (num_tokens, hidden_size)    same as input
+        ln_weight:   (hidden_size)                same as input
+        ln_bias:     (hidden_size)                same as input
+        epsilon:                                  float
+        out:         (num_tokens, hidden_size)    same as input
+    Returns:
+        out:         (num_tokens, hidden_size)    same as input
+    """
+    if out is None:
+        out = torch.empty_like(input)
+    ops.infer.bert_add_norm(input, residual, ln_weight, ln_bias, out, epsilon)
+    return out
+
+
+def ref_bert_unpack_start_end_logits(
+    logits: torch.Tensor,
+    cu_seq_lens: torch.Tensor,
+    max_seq_len: int,
+    start_logits: torch.Tensor = None,
+    end_logits: torch.Tensor = None,
+):
+    batch_size = cu_seq_lens.shape[0] - 1
+    if start_logits is None:
+        start_logits = logits.new_empty([batch_size, max_seq_len])
+    if end_logits is None:
+        end_logits = logits.new_empty([batch_size, max_seq_len])
+    cu_seq_len_cpu = cu_seq_lens.detach().cpu()
+    for i in range(batch_size):
+        start_idx = cu_seq_len_cpu[i]
+        end_idx = cu_seq_len_cpu[i + 1]
+        cur_len = end_idx - start_idx
+        start_logits[i, :cur_len] = logits[start_idx:end_idx, 0]
+        end_logits[i, :cur_len] = logits[start_idx:end_idx, 1]
+    return start_logits, end_logits
+
+
+def bert_unpack_start_end_logits(
+    logits: torch.Tensor,
+    cu_seq_lens: torch.Tensor,
+    max_seq_len: int,
+    start_logits: torch.Tensor = None,
+    end_logits: torch.Tensor = None,
+):
+    """
+    Args:
+        logits:        (num_tokens, 2)              torch.float16, torch.bfloat16
+        cu_seq_lens:    (batch_size+1)               torch.int32, torch.int64
+        max_seq_len:                                int
+        start_logits:  (batch_size, max_seq_len)    same as logits
+        end_logits:    (batch_size, max_seq_len)    same as logits
+    Returns:
+        start_logits:  (batch_size, max_seq_len)    same as logits
+        end_logits:    (batch_size, max_seq_len)    same as logits
+    """
+    batch_size = cu_seq_lens.shape[0] - 1
+    if start_logits is None:
+        start_logits = logits.new_empty([batch_size, max_seq_len])
+    if end_logits is None:
+        end_logits = logits.new_empty([batch_size, max_seq_len])
+    ops.infer.bert_unpack_start_end_logits(
+        logits, cu_seq_lens, start_logits, end_logits
+    )
+    return start_logits, end_logits
+
+
+def ref_bert_linear_residual(
+    input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+):
+    return torch.nn.functional.linear(input, weight, bias) + out
+
+
+def bert_linear_residual(
+    input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+):
+    """
+    Args:
+        input:     (m, k)              torch.float16, torch.bfloat16
+        weight:    (n, k)              same as input
+        bias:      (n)                 same as input
+        out:       (m, n)              same as input
+    Returns:
+        out:       (m, n)              same as input
+    """
+    ops.infer.bert_linear_residual(input, weight, bias, out)
+    return out

ixformer_sdk/inference/functions/bnb_dequant.py

@@ -0,0 +1,55 @@
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+from torch.autograd.function import Function, FunctionCtx
+
+__all__ = [
+    "bnb_dequant",
+    "ref_bnb_dequant",
+]
+
+
+def ref_bnb_dequant(
+    qA: torch.Tensor,
+    SA: torch.Tensor,
+    training: bool = False,
+    scale: float = 127.0,
+    dequant_type: int = 0,
+):
+    A = torch.empty(qA.shape, dtype = SA.dtype, device = SA.device)
+    if dequant_type == 0:
+        for i in range(qA.size(0)): 
+            A[i:] = qA[i:] * (SA[i].to(torch.float) / scale).to(SA.dtype)
+    else:
+        for i in range(qA.size(1)):
+            A[:,i] =  qA[:,i] * (SA[i].to(torch.float) / scale).to(SA.dtype)
+            
+    return A
+
+
+def bnb_dequant(
+    qA: torch.Tensor,
+    SA: torch.Tensor,
+    training: bool = False,
+    scale: float = 127.0,
+    dequant_type: int = 0,
+) -> torch.Tensor:
+    
+    """
+    Args:
+        qA:                  (row, col)            torch.int8
+                dequant input
+        SA:                  (row) or (col)        torch.half
+                scale vector
+        training:                                  bool
+        scale:                                     float
+        dequnt_type:                               int
+                0 : every row shared a scale, SA shape : [row]
+                1 : every col shared a scale, SA shape : [col]
+    Returns:
+        Tensor:              (row, col)            torch.half
+                dequant output    
+        
+    """
+    return ops.infer.bnb_dequant(qA, SA, scale, dequant_type)

ixformer_sdk/inference/functions/bnb_double_quant.py

@@ -0,0 +1,175 @@
+from typing import List, Union
+
+import ixformer._C as ops
+from torch.autograd.function import Function, FunctionCtx
+
+__all__ = ["bnb_double_quant"]
+
+import ctypes as ct
+
+import torch
+from torch import Tensor
+
+
+def get_ptr(A):
+    if A is None:
+        return None
+    else:
+        return ct.c_void_p(A.data.data_ptr())
+
+
+class COOSparseTensor:
+    def __init__(self, rows, cols, nnz, rowidx, colidx, values):
+        assert rowidx.dtype == torch.int
+        assert colidx.dtype == torch.int
+        assert values.dtype == torch.half
+        assert values.numel() == nnz
+        assert rowidx.numel() == nnz
+        assert colidx.numel() == nnz
+
+        self.rows = rows
+        self.cols = cols
+        self.nnz = nnz
+        self.rowidx = rowidx
+        self.colidx = colidx
+        self.values = values
+
+
+def coo_zeros(rows, cols, nnz, device, dtype=torch.half):
+    rowidx = torch.full(size=(nnz,), fill_value=0, dtype=torch.int, device=device)
+
+    colidx = torch.full((nnz,), fill_value=0, dtype=torch.int, device=device)
+    values = torch.full((nnz,), fill_value=0, dtype=dtype, device=device)
+    return COOSparseTensor(rows, cols, nnz, rowidx, colidx, values)
+
+
+def get_colrow_absmax(
+    A, row_stats=None, col_stats=None, nnz_block_ptr=None, threshold=0.0
+):
+    cols = A.shape[-1]
+    if len(A.shape) == 3:
+        rows = A.shape[0] * A.shape[1]
+    else:
+        rows = A.shape[0]
+
+    col_tiles = (cols + 255) // 256
+    tiled_rows = ((rows + 15) // 16) * 16
+    if row_stats is None:
+        row_stats = torch.full(
+            size=(rows,), fill_value=-50000.0, dtype=torch.float, device=A.device
+        )
+    if col_stats is None:
+        col_stats = torch.full(
+            size=(cols,), fill_value=-50000.0, dtype=torch.float, device=A.device
+        )
+
+    # if nnz_block_ptr is None and threshold > 0.0:
+    nnz_block_ptr = torch.full(
+        size=(tiled_rows * col_tiles + 1,),
+        fill_value=0,
+        dtype=torch.int,
+        device=A.device,
+    )
+
+    ops.infer.bnb_getColRowStats(
+        A, row_stats, col_stats, nnz_block_ptr, threshold, rows, cols
+    )
+
+    return row_stats, col_stats, nnz_block_ptr
+
+
+# A : quant input              shape : [row, col] shape:torch.half
+def bnb_double_quant(
+    A: torch.Tensor, training: bool = False, threshold: float = 0.0
+) -> torch.Tensor:
+    
+    """
+    Args:
+        A:                  (row, col)            torch.float16
+                quant input
+        training:                                 bool
+        threshold:                                float
+                abs of element exceeds threshold will be ignored
+    Returns:
+        out_row:            (row, col)            torch.int8
+        out_col:            (row, col)            torch.int8
+        row_stats           (row)                 torch.float
+        col_stats           (col)                 torch.float
+        coo_tensor
+    """
+
+    assert A.dtype == torch.half
+
+    cols = A.shape[-1]
+    if len(A.shape) == 3:
+        rows = A.shape[0] * A.shape[1]
+    else:
+        rows = A.shape[0]
+
+    row_stats, col_stats, nnz_row_ptr = get_colrow_absmax(A, threshold=threshold)
+
+    out_col = torch.full(size=A.shape, fill_value=0, dtype=torch.int8, device=A.device)
+    out_row = torch.full(size=A.shape, fill_value=0, dtype=torch.int8, device=A.device)
+
+    coo_tensor = None
+    if threshold > 0.0:
+        nnz = nnz_row_ptr.cpu().numpy()[-1]
+        if nnz > 0:
+            coo_tensor = coo_zeros(A.shape[0], A.shape[1], nnz, A.device)
+
+            ops.infer.bnb_doubleRowColQuant(
+                A,
+                row_stats,
+                col_stats,
+                out_col,
+                out_row,
+                coo_tensor.rowidx,
+                coo_tensor.colidx,
+                coo_tensor.values,
+                nnz_row_ptr,
+                threshold,
+                rows,
+                cols,
+            )
+            val, idx = torch.sort(torch.Tensor(coo_tensor.rowidx.cpu().numpy()))
+            coo_tensor.rowidx = val
+            coo_tensor.colidx = torch.Tensor(coo_tensor.colidx.cpu().numpy())[idx].to(
+                torch.int32
+            )
+            coo_tensor.values = torch.Tensor(coo_tensor.values.cpu().numpy())[idx].to(
+                torch.half
+            )
+            # coo_tensor.colidx = coo_tensor.colidx[idx]
+            # coo_tensor.values = coo_tensor.values[idx]
+        else:
+            ops.infer.bnb_doubleRowColQuant(
+                A,
+                row_stats,
+                col_stats,
+                out_col,
+                out_row,
+                out_row,
+                out_row,
+                out_row,
+                out_row,
+                0.0,
+                rows,
+                cols,
+            )
+    else:
+        ops.infer.bnb_doubleRowColQuant(
+            A,
+            row_stats,
+            col_stats,
+            out_col,
+            out_row,
+            out_row,
+            out_row,
+            out_row,
+            out_row,
+            threshold,
+            rows,
+            cols,
+        )
+
+    return out_row, out_col, row_stats, col_stats, coo_tensor

ixformer_sdk/inference/functions/bnb_mm_dequant.py

@@ -0,0 +1,73 @@
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+from torch.autograd.function import Function, FunctionCtx
+
+__all__ = ["bnb_mm_dequant"]
+
+
+# A : quant input              shape : [row, col]  shape : torch.int
+def bnb_mm_dequant(
+    A: torch.Tensor,
+    quant_state: tuple,
+    row_stats: torch.Tensor,
+    col_stats: torch.Tensor,
+    bias: torch.Tensor = None,
+    add_bias: bool = False,
+    training: bool = False,
+) -> torch.Tensor:
+    
+    """
+    Args:
+        A:                  (row, col)            torch.int8
+        quant_state:                              tuple
+        row_stats:          (row)                 torch.float
+        col_stats:          (col)                 torch.float
+        bias:               (col)                 torch.half
+        add_bias:                                 bool
+        training:                                 bool
+    Returns:
+        Tensor:             (row, col)            torch.half
+        
+    """
+    
+    assert A.dtype == torch.int
+    if bias is not None:
+        add_bias = True
+        print("bias.dtype:", bias.dtype)
+        assert bias.dtype == torch.half
+    else:
+        bias = A
+    out_shape = quant_state[0]
+    if len(out_shape) == 3:
+        out_shape = (out_shape[0] * out_shape[1], out_shape[2])
+    out = torch.full(size=out_shape, fill_value=0, dtype=torch.half, device=A.device)
+    new_row_stats = torch.full(
+        size=(out_shape[0],), fill_value=0, dtype=torch.float, device=A.device
+    )
+    new_col_stats = torch.full(
+        size=(out_shape[1],), fill_value=0, dtype=torch.float, device=A.device
+    )
+
+    assert (
+        new_row_stats.shape[0] == row_stats.shape[0]
+    ), f"{new_row_stats.shape} vs {row_stats.shape}"
+    assert (
+        new_col_stats.shape[0] == col_stats.shape[0]
+    ), f"{new_col_stats.shape} vs {col_stats.shape}"
+    numRows = out_shape[0]
+    numCols = out_shape[1]
+    ops.infer.bnb_mm_dequant(
+        A,
+        row_stats,
+        col_stats,
+        out,
+        new_row_stats,
+        new_col_stats,
+        numRows,
+        numCols,
+        add_bias,
+        bias,
+    )
+    return out

ixformer_sdk/inference/functions/bnb_qgemm.py

@@ -0,0 +1,56 @@
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+
+__all__ = ["bnb_qgemm", "ref_bnb_qgemm"]
+
+
+# qA : quant input              shape : [bs, in_feature]
+# qW : quant weight             shape : [out_feature, in_feature]
+# SA : scale vector of qA       shape : [bs]
+# SW : scale vector of qW       shape : [out_feature]
+
+
+def ref_bnb_qgemm(
+    qA: torch.Tensor,
+    qW: torch.Tensor,
+    SA: torch.Tensor,
+    SW: torch.Tensor,
+    training: bool = False,
+    scaleA: float = 127.0,
+    scaleW: float = 127.0,
+):
+    y = torch.nn.functional.linear(qA.to(torch.float), qW.to(torch.float))
+    out = torch.empty(y.shape, dtype = SA.dtype, device = SA.device)
+    for i in range(qA.size(0)):
+        for j in range(qW.size(0)): 
+            out[i][j] = y[i][j] * (SA[i].to(torch.float) / scaleA) * (SW[j].to(torch.float) / scaleW)
+    return out.to(SA.dtype)
+
+
+def bnb_qgemm(
+    qA: torch.Tensor,
+    qW: torch.Tensor,
+    SA: torch.Tensor,
+    SW: torch.Tensor,
+    training: bool = False,
+    scaleA: float = 127.0,
+    scaleW: float = 127.0,
+) -> torch.Tensor:
+    
+    """
+    Args:
+        qA:                 (bs, in_feature)                torch.int8
+        qW:                 (out_feature, in_feature)       torch.int8
+        SA:                 (bs)                            torch.half
+            scale vector of qA 
+        SA:                 (out_feature)                   torch.half
+            scale vector of qW 
+        training:                                           bool
+        scaleA:                                             float
+        scaleW:                                             float
+    Returns:
+        Tensor:             (bs, out_feature)               torch.half
+    """
+    return ops.infer.bnb_qgemm(qA, qW, SA, SW, scaleA, scaleW)

ixformer_sdk/inference/functions/bnb_quant.py

@@ -0,0 +1,58 @@
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+
+__all__ = ["bnb_quant", "ref_bnb_quant"]
+
+
+# A : input   shape : [row, col]
+# SA : scale vector
+# quant_type
+#   0 : every row shared a scale, SA shape : [row]
+#   1 : every col shared a scale, SA shape : [col]
+def ref_bnb_quant(
+    A: torch.Tensor,
+    SA: torch.Tensor,
+    training: bool = False,
+    scale: float = 127.0,
+    quant_type: int = 0,
+):
+    qA = torch.empty(A.shape, device = SA.device)
+    if quant_type == 0:
+        for i in range(A.size(0)):
+            qA[i:] = torch.round(A[i:] * (scale / SA[i].to(torch.float)))
+    else:
+        for i in range(A.size(1)):
+            qA[:,i] =  torch.round(A[:,i] * (scale / SA[i].to(torch.float)))
+
+    qA_clamped = torch.clamp(qA, min=-128, max=127)
+    qA = qA_clamped.to(torch.int8)
+    return qA
+
+
+def bnb_quant(
+    A: torch.Tensor,
+    SA: torch.Tensor,
+    training: bool = False,
+    scale: float = 127.0,
+    quant_type: int = 0,
+) -> torch.Tensor:
+    
+    """
+    Args:
+        A:                   (row, col)            torch.half
+                quant input
+        SA:                  (row) or (col)        torch.half
+                scale vector
+        training:                                  bool
+        scale:                                     float
+        qunt_type:                                 int
+                0 : every row shared a scale, SA shape : [row]
+                1 : every col shared a scale, SA shape : [col]
+    Returns:
+        Tensor:              (row, col)            torch.int8
+                quant output    
+        
+    """
+    return ops.infer.bnb_quant(A, SA, scale, quant_type)

ixformer_sdk/inference/functions/bnb_rowcol_absmax.py

@@ -0,0 +1,52 @@
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+
+__all__ = ["bnb_rowcol_absmax", "ref_bnb_rowcol_absmax"]
+
+
+# input  : input   shape : [row, col]
+# threshold : abs of element exceeds threshold will be ignored
+# type
+#   0 : row absmax
+def ref_bnb_rowcol_absmax(
+    input: torch.Tensor,
+    training: bool = False,
+    threshold: float = 0.0,
+    type: int = 0,
+):
+    input = input.float()
+    if threshold ==0.0:
+        threshold = float('inf')
+    mask = (torch.abs(input) < threshold)
+    masked_input = mask * input
+    masked_input = masked_input.half()
+    if type == 0:
+        out = torch.amax(torch.abs(masked_input), dim=1)
+        
+    else:
+        out = torch.amax(torch.abs(masked_input), dim=0)
+    return out
+
+
+def bnb_rowcol_absmax(
+    input: torch.Tensor,
+    training: bool = False,
+    threshold: float = 0.0,
+    type: int = 0,
+) -> torch.Tensor:
+    
+    """
+    Args:
+        input:              (row, col)        torch.half
+                目前col值必须满足col%2==0
+        training:                             bool
+        threshold:                            float
+                abs of element exceeds threshold will be ignored
+        type:                                  int
+                row absmax, 目前只支持type=0
+    Returns:
+        Tensor:             (row)              torch.half
+    """
+    return ops.infer.bnb_rowcol_absmax(input, threshold, type)

ixformer_sdk/inference/functions/conv2d.py

@@ -0,0 +1,198 @@
+from typing import Union
+
+import ixformer._C as ops
+import torch
+import torch.nn.functional as NNF
+
+__all__ = ["conv2d", "ref_conv2d", "ref_conv2d_nhwc", "conv2d_nhwc"]
+
+
+def is_channels_last(ten):
+    return torch._prims_common.suggest_memory_format(ten) == torch.channels_last
+
+
+def _pair(x):
+    if isinstance(x, (list, tuple)):
+        return x
+    return (x, x)
+
+
+def ref_conv2d(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor = None,
+    stride: Union[int, tuple] = 1,
+    padding: Union[int, tuple] = 0,
+    dilation: Union[int, tuple] = 1,
+    groups: int = 1,
+):
+
+    output = NNF.conv2d(input, weight, bias, stride, padding, dilation, groups)
+    return output
+
+
+# conv2d官方接口，如果weight是torch.channels_last,输出也是torch.channels_last；如果weight是nchw，那么输出也是nchw;特殊情况，如果输入是nchw，weight是torch.channels_last，输出也是torch.channels_last
+def conv2d(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor = None,
+    stride: Union[int, tuple] = 1,
+    padding: Union[int, tuple] = 0,
+    dilation: Union[int, tuple] = 1,
+    groups: int = 1,
+):
+    """
+    Args:
+        input:              (n,in_c,h,w)                        torch.float16
+        weight:             (out_c,in_c/groups,kH,kW)           torch.float16
+        bias:               (out_c)                             torch.float16
+        stride:                                                 int or tuple
+                Stride of the convolution. Default: 1
+        padding:                                                int or tuple
+                Padding added to all four sides of the input. Default: 0
+        dilation:                                               int or tuple
+                Spacing between kernel elements. Default: 1
+        groups:                                                 int
+                Number of blocked connections from input channels to output channels. Default: 1
+    Returns:
+        Tensor:             (n,out_c,h_out,w_out)               torch.float16
+                h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1;
+                w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1;      
+    """
+    stride = _pair(stride)
+    padding = _pair(padding)
+    dilation = _pair(dilation)
+
+    channel_last = is_channels_last(weight)
+    if not is_channels_last(input) and channel_last:
+        input = input.to(memory_format=torch.channels_last)
+
+    # compute outshape
+    n, in_c, h_in, w_in = input.shape
+    out_c, _, kernel_h, kernel_w = weight.shape
+    pad_h = padding[0]
+    pad_w = padding[1]
+    stride_h = stride[0]
+    stride_w = stride[1]
+    dilation_h = dilation[0]
+    dilation_w = dilation[1]
+    h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1
+    w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1
+
+    if channel_last:
+        output_shape = [n, out_c, h_out, w_out]
+        output = torch.empty(
+            output_shape,
+            memory_format=torch.channels_last,
+            dtype=input.dtype,
+            device=input.device,
+        )
+    else:
+        output_shape = [n, out_c, h_out, w_out]
+        output = input.new_empty(output_shape)
+
+    if channel_last:
+        input = input.permute(0, 2, 3, 1)
+        weight = weight.permute(0, 2, 3, 1)
+        output = output.permute(0, 2, 3, 1)
+
+    if bias is not None:
+        bias = bias.float()
+    ops.infer.conv2d(
+        input, weight, bias, output, stride, padding, dilation, groups, channel_last
+    )
+    if channel_last:
+        output = output.permute(0, 3, 1, 2)
+
+    return output
+
+
+def ref_conv2d_nhwc(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor = None,
+    stride: Union[int, tuple] = 1,
+    padding: Union[int, tuple] = 0,
+    dilation: Union[int, tuple] = 1,
+    groups: int = 1,
+):
+
+    output = NNF.conv2d(
+        input.permute(0, 3, 1, 2).contiguous(),
+        weight.permute(0, 3, 1, 2).contiguous(),
+        bias,
+        stride,
+        padding,
+        dilation,
+        groups,
+    )
+    return output.permute(0, 2, 3, 1).contiguous()
+
+
+# conv2d_nhwc，
+# conv2d官方接口解决两种情况：
+# 1、务必输入tensor内存上是nhwc,且tensor属于memory_format=torch.channels_last，
+# 2、或者输入tensor内存上是nchw，并且是contiguous；
+# conv2d官方接口不能解决，conv2d_nhwc则可处理这种情况的
+# 输入tensor内存上是nhwc的，但tensor没有用memory_format=torch.channels_last进行过处理，不会有memory_format=torch.channels_last的标签
+def conv2d_nhwc(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor = None,
+    stride: Union[int, tuple] = 1,
+    padding: Union[int, tuple] = 0,
+    dilation: Union[int, tuple] = 1,
+    groups: int = 1,
+):
+    
+    """
+    Args:
+        input:              (n,h,w,in_c)                        torch.float16
+        weight:             (out_c,kH,kW,in_c/groups)           torch.float16
+        bias:               (out_c)                             torch.float16
+        stride:                                                 int or tuple
+                Stride of the convolution. Default: 1
+        padding:                                                int or tuple
+                Padding added to all four sides of the input. Default: 0
+        dilation:                                               int or tuple
+                Spacing between kernel elements. Default: 1
+        groups:                                                 int
+                Number of blocked connections from input channels to output channels. Default: 1
+    Returns:
+        Tensor:             (n,h_out,w_out,out_c)               torch.float16
+                h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1;
+                w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1;      
+    """
+
+    stride = _pair(stride)
+    padding = _pair(padding)
+    dilation = _pair(dilation)
+
+    assert input.is_contiguous()
+    assert weight.is_contiguous()
+
+    # compute outshape
+    n, h_in, w_in, in_c = input.shape
+    (
+        out_c,
+        kernel_h,
+        kernel_w,
+        _,
+    ) = weight.shape
+    pad_h = padding[0]
+    pad_w = padding[1]
+    stride_h = stride[0]
+    stride_w = stride[1]
+    dilation_h = dilation[0]
+    dilation_w = dilation[1]
+    h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1
+    w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1
+
+    output_shape = [n, h_out, w_out, out_c]
+    output = torch.empty(output_shape, dtype=input.dtype, device=input.device)
+    if bias is not None:
+        bias = bias.float()
+    ops.infer.conv2d(
+        input, weight, bias, output, stride, padding, dilation, groups, True
+    )
+    return output

ixformer_sdk/inference/functions/cross_entropy_loss.py

@@ -0,0 +1,204 @@
+from typing import Union
+
+import ixformer._C as ops
+import torch
+
+__all__ = ["vocab_parallel_cross_entropy", "ref_vocab_parallel_cross_entropy"]
+
+
+def ref_vocab_parallel_cross_entropy(
+    vocab_parallel_logits: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float = 0.0,
+    world_size: int = 1,
+    vocab_start_index: int = 0,
+    vocab_end_index: int = 320000,
+    group=None,
+):
+    if world_size == 1:
+        vocab_parallel_logits = vocab_parallel_logits.float()
+        partition_vocab_size = vocab_parallel_logits.size()[-1]
+        logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
+        vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1)
+        target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
+        masked_target = target.clone() - vocab_start_index
+        masked_target[target_mask] = 0
+        logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
+        masked_target_1d = masked_target.view(-1)
+        arange_1d = torch.arange(
+            start=0, end=logits_2d.size()[0], device=logits_2d.device
+        )
+
+        predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
+        predicted_logits_1d = predicted_logits_1d.clone().contiguous()
+
+        predicted_logits = predicted_logits_1d.view_as(target)
+        predicted_logits[target_mask] = 0.0
+        exp_logits = vocab_parallel_logits
+        torch.exp(vocab_parallel_logits, out=exp_logits)
+        sum_exp_logits = exp_logits.sum(dim=-1)
+        loss = torch.log(sum_exp_logits) - predicted_logits
+        exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
+        if label_smoothing > 0:
+            """
+            We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
+            = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
+            = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
+            = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
+            = (K * (1 - alpha) - 1) / (K - 1)) * y_gt  + (alpha / (K - 1)) * \sum_{i} y_i
+            = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
+            From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
+            """
+            assert 1.0 > label_smoothing > 0.0
+            smoothing = label_smoothing * partition_vocab_size / (partition_vocab_size - 1)
+
+            # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
+            log_probs = torch.log(exp_logits)
+            mean_log_probs = log_probs.mean(dim=-1)
+            loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
+
+    else:
+        # Maximum value along vocab dimension across all GPUs.
+        logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
+        torch.distributed.all_reduce(
+            logits_max, op=torch.distributed.ReduceOp.MAX, group=group
+        )
+        # Subtract the maximum value.
+        vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1)
+
+        # Get the partition's vocab indecies
+        partition_vocab_size = vocab_parallel_logits.size()[-1]
+
+        # Create a mask of valid vocab ids (1 means it needs to be masked).
+        target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
+        masked_target = target.clone() - vocab_start_index
+        masked_target[target_mask] = 0
+
+        # Get predicted-logits = logits[target].
+        # For Simplicity, we convert logits to a 2-D tensor with size
+        # [*, partition-vocab-size] and target to a 1-D tensor of size [*].
+        logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
+        masked_target_1d = masked_target.view(-1)
+        arange_1d = torch.arange(
+            start=0, end=logits_2d.size()[0], device=logits_2d.device
+        )
+        predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
+        predicted_logits_1d = predicted_logits_1d.clone().contiguous()
+        predicted_logits = predicted_logits_1d.view_as(target)
+        predicted_logits[target_mask] = 0.0
+        # All reduce is needed to get the chunks from other GPUs.
+        torch.distributed.all_reduce(
+            predicted_logits,
+            op=torch.distributed.ReduceOp.SUM,
+            group=group,
+        )
+
+        # Sum of exponential of logits along vocab dimension across all GPUs.
+        exp_logits = vocab_parallel_logits
+        torch.exp(vocab_parallel_logits, out=exp_logits)
+        sum_exp_logits = exp_logits.sum(dim=-1)
+        torch.distributed.all_reduce(
+            sum_exp_logits,
+            op=torch.distributed.ReduceOp.SUM,
+            group=group,
+        )
+
+        # Loss = log(sum(exp(logits))) - predicted-logit.
+        loss = torch.log(sum_exp_logits) - predicted_logits
+
+        # Normalize and optionally smooth logits
+        exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
+
+        vocab_size = exp_logits.size(-1)
+        if label_smoothing > 0:
+            """
+            We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
+            = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
+            = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
+            = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
+            = (K * (1 - alpha) - 1) / (K - 1)) * y_gt  + (alpha / (K - 1)) * \sum_{i} y_i
+            = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
+            From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
+            """
+            assert 1.0 > label_smoothing > 0.0
+            smoothing = label_smoothing * vocab_size / (vocab_size - 1)
+
+            # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
+            log_probs = torch.log(exp_logits)
+            mean_log_probs = log_probs.mean(dim=-1)
+            loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
+
+    return loss
+
+
+def vocab_parallel_cross_entropy(
+    vocab_parallel_logits: torch.Tensor,
+    target: torch.Tensor,
+    label_smoothing: float = 0.0,
+    world_size: int = 1,
+    vocab_start_index: int = 0,
+    vocab_end_index: int = 320000,
+    group=None,
+):
+    """
+    Args:
+        vocab_parallel_logits:              (seq_len,1,vocal_size)              torch.float16, torch.bfloat16, torch.float
+        target:                             (seq_len,1)                         torch.int64
+        label_smoothing:                                                        float
+                默认为0.0. 用于标签平滑
+        world_size:                                                             int
+                当world_size = 1时,目前只支持batch_size = 1 的情况
+        vocab_start_index:                                                      int
+        vocab_end_index:                                                        int
+        group:  
+                TP 并行组                                                  
+    Returns:
+        loss:                                (seq_len,1)                        torch.float          
+    """
+    if world_size == 1:
+        device = vocab_parallel_logits.device
+        xnumel = vocab_parallel_logits.shape[0]
+        rnumel = vocab_parallel_logits.shape[-1]
+        exp_logits = torch.empty(
+            (xnumel, 1, rnumel), device=device, dtype=torch.float32
+        )
+        masked_target_1d = torch.empty((xnumel,), device=device, dtype=torch.int32)
+        loss = torch.empty((xnumel, 1), device=device, dtype=torch.float32)
+
+        ops.train.cross_entropy_loss_forward(
+            vocab_parallel_logits, target.int(), exp_logits, masked_target_1d, loss
+        )
+
+        vocab_size = exp_logits.size(-1)
+        if label_smoothing > 0:
+            """
+            We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
+            = (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
+            = (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
+            = ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
+            = (K * (1 - alpha) - 1) / (K - 1)) * y_gt  + (alpha / (K - 1)) * \sum_{i} y_i
+            = (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
+            From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
+            """
+            assert 1.0 > label_smoothing > 0.0
+            smoothing = label_smoothing * vocab_size / (vocab_size - 1)
+
+            # Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
+            log_probs = torch.log(exp_logits)
+            mean_log_probs = log_probs.mean(dim=-1)
+            loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
+
+        # Store softmax, target-mask and masked-target for backward pass.
+
+        return loss
+    else:
+        loss = ref_vocab_parallel_cross_entropy(
+            vocab_parallel_logits,
+            target,
+            label_smoothing,
+            world_size,
+            vocab_start_index,
+            vocab_end_index,
+            group,
+        )
+        return loss

ixformer_sdk/inference/functions/flash_attn.py

@@ -0,0 +1,201 @@
+import math
+from typing import List, Union
+
+import ixformer._C as ops
+import torch
+from torch.autograd.function import Function, FunctionCtx
+
+__all__ = [
+    "ixinfer_flash_attn_unpad",
+    "ixinfer_flash_attn_pad",
+    "ref_ixinfer_flash_attn_pad",
+]
+
+
+def ixinfer_flash_attn_unpad(
+    # total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+    q: "torch.Tensor",
+    # total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
+    k: "torch.Tensor",
+    # total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
+    v: "torch.Tensor",
+    # total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    cu_seqlens_q: "torch.Tensor",  # b+1
+    cu_seqlens_k: "torch.Tensor",  # b+1
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    is_causal: bool = False,
+    atten_scale: float = None,
+    sqrt_alibi: bool = False,
+    # total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+    alibi_slopes: "torch.Tensor" = None,
+    out: "torch.Tenosr" = None,
+):
+    """
+    Args:
+        q:              (total_q, nheads, headdim)          torch.float16, torch.bfloat16
+            where total_q = total number of query tokens in the batch. 
+        k:              (total_k, nheads_k, headdim)        torch.float16, torch.bfloat16
+            where total_k = total number of key tokens in the batch.
+        v:              (total_k, nheads_k, headdim)        torch.float16, torch.bfloat16   
+        cu_seqlens_q:   (batch_size + 1)                    torch.int32
+            The cumulative sequence lengths of the sequences in the batch, used to index into q.
+        cu_seqlens_k:   (batch_size + 1)                    torch.int32
+            The cumulative sequence lengths of the sequences in the batch, used to index into kv.
+        max_seqlen_q:                                       int
+            Maximum query sequence length in the batch.
+        max_seqlen_k:                                       int
+            Maximum key sequence length in the batch.
+        atten_scale:                                        float
+            The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim).
+        is_causal:                                          bool
+            Whether to apply causal attention mask (e.g., for auto-regressive modeling).
+         sqrt_alibi:                                        bool
+            Whether to apply abilimode
+        out:            (total, nheads, headdim)            torch.float16, torch.bfloat16
+    Returns:
+        out:            (total, nheads, headdim)            torch.float16, torch.bfloat16
+            if not q.size(-1) % 32 == 0: out shape is (total_q, nheads, q.size(-1) + (32 - q.size(-1) % 32))
+    """
+    if atten_scale is None:
+        atten_scale = 1.0 / (q.size(-1) ** 0.5)
+
+    # 判断是否pad
+    cur_head = q.size(-1)
+    cur_head32 = cur_head
+    if not cur_head % 32 == 0:
+        cur_head32 = cur_head + (32 - cur_head % 32)
+        q_infer = torch.nn.functional.pad(q, [0, cur_head32 - cur_head, 0, 0], value=0)
+        k_infer = torch.nn.functional.pad(k, [0, cur_head32 - cur_head, 0, 0], value=0)
+        v_infer = torch.nn.functional.pad(v, [0, cur_head32 - cur_head, 0, 0], value=0)
+    else:
+        q_infer = q
+        k_infer = k
+        v_infer = v
+
+    if out is None:
+        out = torch.empty_like(q_infer)
+    # ixinfer 新接口版
+    ops.infer.ixinfer_flash_attn_unpad(
+        q_infer,
+        k_infer,
+        v_infer,
+        out,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        is_causal,
+        False,  # need_lse =False
+        atten_scale,
+        sqrt_alibi,
+        alibi_slopes,
+    )
+    if not cur_head % 32 == 0:
+        out = out[:, :, :cur_head]
+    return out
+
+
+def ref_ixinfer_flash_attn_pad(
+    # [ batch  num_heads seq_q head_size]
+    q: torch.Tensor,
+    # [ batch  num_heads_k max_seq_kv head_size]
+    k: torch.Tensor,
+    # [ batch  num_heads_k max_seq_kv head_size]
+    v: torch.Tensor,
+    # [ batch  num_heads seq_q seq_kv] seq_kv<=max_seq_kv
+    mask: torch.Tensor,
+    # [ batch  num_heads seq_q head_size]
+    atten_scale: float = None,
+    kv_seq_start: int = None,
+    kv_seq_end: int = None,
+):
+    head_dim = q.size(-1)
+    k_effective = k[:, :, kv_seq_start:kv_seq_end, :]
+    v_effective = v[:, :, kv_seq_start:kv_seq_end, :]
+    # 2. q*kt softmax
+    scores_qk = (
+        torch.matmul(q.float(), k_effective.float().transpose(-2, -1)) * atten_scale
+    )
+    # softmax
+    # print(scores_qk.shape,mask.shape)
+    if mask is not None:
+        if mask.dtype == torch.int32:
+            scores_qk = scores_qk + mask * (-100000)
+        elif mask.dtype == torch.float32:
+            scores_qk = scores_qk + mask
+        else:
+            print(
+                f"mask dtype is not surported {mask.dtype},now surport int32 and float32"
+            )
+    scores_qk = torch.nn.functional.softmax(scores_qk, dim=-1)
+    # 3.  x = qk_scores * v
+    scores_v = torch.matmul(scores_qk, v_effective.float())
+    return scores_v.half()
+
+
+def ixinfer_flash_attn_pad(
+    # [ batch  num_heads seq_q head_size]
+    q: torch.Tensor,
+    # [ batch  num_heads_k max_seq_kv head_size]
+    k: torch.Tensor,
+    # [ batch  num_heads_k max_seq_kv head_size]
+    v: torch.Tensor,
+    # [ batch  num_heads seq_q seq_kv] seq_kv<=max_seq_kv
+    mask: torch.Tensor,
+    # [ batch  num_heads seq_q head_size]
+    atten_scale: float = None,
+    kv_seq_start: int = None,
+    kv_seq_end: int = None,
+):
+    """
+    Args:
+        q:              (batch_size, num_head, seq_len_q, head_dim)                     torch.float16, torch.bfloat16
+        k:              (batch_size, num_head_kv, seq_len_kv, head_dim)                 torch.float16, torch.bfloat16
+        v:              (batch_size, num_head_kv, seq_len_kv, head_dim)                 torch.float16, torch.bfloat16   
+        mask:           (batch_size, num_head, seq_len_q, kv_seq_start:kv_seq_end)      torch.int32, torch.int64, torch.float32
+        atten_scale:                                                                    float
+            The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim).
+        kv_seq_start:                                                                   int
+            kv sequence start index used for computation in the batch
+        kv_seq_end:                                                                     int
+            kv sequence end index used for computation in the batch.
+    Returns:
+        out:            (batch_size, num_head, seq_len_q, head_dim)                     torch.float16, torch.bfloat16
+    """
+
+    # 判断是否pad
+    cur_head = q.size(-1)
+    cur_head32 = cur_head
+    if not cur_head % 32 == 0:
+        cur_head32 = cur_head + (32 - cur_head % 32)
+        q_infer = torch.nn.functional.pad(q, [0, cur_head32 - cur_head, 0, 0], value=0)
+        k_infer = torch.nn.functional.pad(k, [0, cur_head32 - cur_head, 0, 0], value=0)
+        v_infer = torch.nn.functional.pad(v, [0, cur_head32 - cur_head, 0, 0], value=0)
+    else:
+        q_infer = q
+        k_infer = k
+        v_infer = v
+
+    if atten_scale is None:
+        atten_scale = 1.0 / (q.size(-1) ** 0.5)
+    if kv_seq_start is None or kv_seq_end is None:
+        kv_seq_start = 0
+        kv_seq_end = k.size(-2)  # kv seq len
+    elif kv_seq_start < 0 or kv_seq_end > k.size(-2) or kv_seq_start >= kv_seq_end:
+        raise NotImplementedError(
+            "must kv_seq_start<0 or kv_seq_end>k.size(-2) or kv_seq_start>=kv_seq_end!"
+        )
+    out_shape = list(q_infer.shape)
+    out = torch.empty(out_shape, dtype=q.dtype, device=q.device)
+    if mask is not None:
+        ops.infer.ixinfer_flash_attn_pad_fwd(
+            q_infer, k_infer, v_infer, mask, out, atten_scale, kv_seq_start, kv_seq_end
+        )
+    else:
+        ops.infer.ixinfer_flash_attn_pad_fwd_nomask(
+            q_infer, k_infer, v_infer, out, atten_scale, kv_seq_start, kv_seq_end
+        )
+    if not cur_head % 32 == 0:
+        out = out[:, :, :, :cur_head]
+    return out