trigger build

- reshape_and_cache → reshape_paged_cache (xllm_cache.so)
- residual_rms_norm → fused_add_rms_norm (xllm_norm.so)
- topk_softmax → moe_fused_topk (xllm_moe.so)
- moe_compute_token_index → moe_compute_index (xllm_moe.so)
- ix_full_bridge → ix_moe_bridge (paged_attention, ix_linear)
- ix_linear 5参数 → 3参数 (input, weight, bias)
- check_all: ix_moe_bridge改为required, ix_full_bridge改为optional

.dockerignore

@@ -0,0 +1,16 @@
+**/__pycache__
+**/*.pyc
+**/.git
+cccl_upstream/
+upstream_ref/
+vllm/
+ixformer_sdk/
+muh/
+ex_engine/fla_kernels/
+ex_engine/moe/
+ex_engine/xllm_layers/npu_torch/
+ex_engine/xllm_layers/mlu/
+ex_engine/xllm_models/
+*.zip
+dockerrizhi.txt
+subrizhi.txt

.gitattributes

@@ -0,0 +1,10 @@
+# Force LF line endings for all text files
+* text=auto eol=lf
+*.py text eol=lf
+*.sh text eol=lf
+*.cu text eol=lf
+*.cuh text eol=lf
+*.yaml text eol=lf
+*.yml text eol=lf
+*.md text eol=lf
+Dockerfile text eol=lf

.gitignore

@@ -6,3 +6,4 @@ baseline.muh
 pkgs/
 enginex_base/
 __pycache__/
+*.pyc

CCCL_INTEGRATION_STATUS.md

@@ -36,7 +36,6 @@
 | schema YAML | muh/schema/*.yaml | 27 files | ✓ 完成 | 每个算法的参数空间定义 |
 | parse.py | muh/parse.py | ~150 | ✓ 基本可用 | .muh → JSON 解析（自实现 YAML parser）|
 | gen_yaml.py | muh/gen_yaml.py | ~80 | ✓ 完成 | .muh → computility-run.yaml |
-| gen_patch.py | muh/gen_patch.py | ~200 | ✓ 可提取 bi100_* | C++ header → vllm unified diff |
 | extract.py | muh/extract.py | ~100 | ✓ 完成 | CCCL tuning → schema 提取 |
 | muh_dispatch.py | muh_dispatch.py | ~400 | △ 概念完成 | CCCL-style 类型分派（未接入 vllm）|
 | muh_kernel_map.py | muh_kernel_map.py | ~350 | △ 手写常量 | 需要从 C++ headers 自动提取闭环 |
@@ -44,98 +43,3 @@
 | test_smem_safety | muh/tests/test_smem_safety.py | 1 file | ✓ | 全算法 SMEM 安全检查 |
 
 ---
-
-## 三、Decode 热路径 × 资产覆盖矩阵
-
-```
-算法              CCCL  muh   schema bench test  vllm注入点                         竞赛权重
-───────────────── ───── ───── ────── ───── ───── ──────────────────────────────── ─────────
-reduce            ✓     ✓     ✓      ✓     ✓     csrc/attention/paged_attention   83% (Output)
-scan              ✓     ✓     ✓      ✓     ✓     csrc/attention/paged_attention   14% (Input)
-topk              ✓     ✓     ✓      ✓     ✓     csrc/sampling/sampling_kernels   per decode
-radix_sort        ✓     ✓     ✓      ✓     ✓     csrc/sampling/sampling_kernels   per decode
-transform         ✓     ✓     ✓      ✓     ✓     csrc/activation/layernorm/rope   200×/token
-select_if         ✓     ✓     ✓      △     ✓     csrc/sampling (top-p filter)     per decode
-batch_memcpy      ✓     ✓     ✓      △     △     csrc/cache_kernels               3% (Cache)
-for_each          ✓     ✓     ✓      ✓     ✓     csrc (residual connections)      per layer
-```
-
-△ = benchmark/test 文件存在但名称不直接匹配（partition/if.cu 对应 select_if，copy/memcpy.cu 对应 batch_memcpy）
-
----
-
-## 四、GitHub Issues 状态
-
-### 已创建的 38 个 Issues（全 open，全有 labels）
-
-**功能测试覆盖（#1-#16）**: 竞赛 50+ 功能测试用例的完整 PRD，每个含 PND 级 test cases 表
-
-| 编号范围 | 前缀 | 数量 | 说明 |
-|----------|------|------|------|
-| #1-#14 | [FEA] | 14 | 功能测试: 非流式/流式/Tool/Reasoning/Cache/采样/结构化/多语言/多模态/校验/能力/截断/效果 |
-| #15-#16 | [EPIC] | 2 | 性能基准 + 开发环境 |
-| #17-#25 | [FEA]/[EPIC] | 9 | muh 语言设计: 语法/schema/codegen(yaml+patch+dockerfile)/bench/search/tuning提取 |
-| #26-#38 | [muh] | 13 | muh 算法标定: reduce/scan/radix_sort/select_if/scan_by_key/reduce_by_key/unique_by_key/transform/batch_memcpy/topk + gen_patch管道/benchmark runner/hardware校准 |
-
-### Project/6 面板上的 Draft Issues（72 个，无 repo 关联）
-
-来自后续对话生成，包含：
-- [muh] 语言规范 v1/v2
-- [muh] 20+ 个算法标定 items（adjacent_difference, batched_topk, find, histogram, merge, rle_encode 等）
-- [INFRA] CI同步/Build编译/Deploy部署/Verify回归
-- [BUG] scale_mem_bound / gen_patch 管道 / select_if 坍缩 / bytes_in_flight / reduce items
-- [CCCL-verify] 20 个 Thrust/CUB example 验证 items
-- [CCCL-test] 10 个 Catch2 测试矩阵 items
-- [muh-bench] 6 个 benchmark items
-- [muh-pipe] 端到端管道验证
-
-**这 72 个 draft 需要转为真 issue。** 内容已经写好（body 含完整 test cases 表），只是缺少 repo 关联和 labels。
-
----
-
-## 五、关键发现（SM count = 16）
-
-Phanthy Cloud 实测确认 BI-V100 只有 **16 SMs**（不是规格书的 50c）。
-
-影响范围：
-1. `hardware.cuh` — 已修正 sm_count=16
-2. `tuning_transform.cuh` — bytes_in_flight 基于 900/50=18 GB/s 已失效，应为 900/16=56 GB/s
-3. `tuning_reduce.cuh` — bi100_det_* 和 bi100_default 的 items 偏小（tile 仅用 23% SMEM）
-4. `tuning_scan.cuh` — lookback delay 基于 50 SM 的争用模型，16 SM 下争用更低、delay 可以更短
-5. 所有 benchmark 理论推导需要重跑
-
----
-
-## 六、不需要 clone 更多 CCCL 的原因
-
-完整 CCCL (github.com/NVIDIA/cccl) ≈ 40K 文件、1.2GB。我们有 8,900 文件 (74MB)。
-
-已有的关键子集：
-- ✓ 全部 27 tuning headers（muh 从这里提取参数空间）
-- ✓ 全部 32 dispatch headers（tuning 参数化的对象）
-- ✓ 52 Thrust examples（正确性验证的 golden reference）
-- ✓ 18 CUB examples（device + block level API 验证）
-- ✓ 234 CUB Catch2 tests（回归测试矩阵）
-- ✓ 169 Thrust tests（Thrust 算法回归）
-- ✓ 78 CUB benchmarks（标定数据的来源）
-- ✓ 530 Thrust headers + 1,357 libcudacxx headers（编译依赖）
-
-缺失的 ~31K 文件：
-- libcudacxx 深层 include（6K）— 编译时用 -I 指向安装路径
-- cudax 实验模块（800）— 竞赛不用
-- cmake/CI 基础设施（5K）— 平台用 Dockerfile 构建
-- Python 绑定 / 文档 / 其他（19K）— 不相关
-
----
-
-## 七、信创魔盒核心差异（竞赛定位）
-
-> "信创魔盒是基于系统级的架构，内置算法因子，用 EngineX 引擎把模型内部的算法因子重新置换——不是单纯的连接器。"
-
-muh 在这个架构中的角色：
-- CCCL 的 `policy_selector` 是 NVIDIA 为自家 GPU 写的"算法因子"
-- muh 的 `policy_selector` 是为天垓100 写的等效"算法因子"
-- EngineX 把 CCCL 的 NVIDIA 算法因子替换成 muh 的天垓100 算法因子
-- 不是适配层（60% 精度），是置换层（目标 ≥100% 精度在天垓100 硬件约束下的最优解）
-
-竞赛成绩 = 算法因子置换的精度 × 硬件实测标定的覆盖度。

CCCL_MUH_GAP_ANALYSIS.md

@@ -1,120 +0,0 @@
-# CCCL ↔ muh 完整 Gap 分析
-
-> 生成时间: 2026-08-06 | HEAD: 2a7ca10 | 26 算法全量扫描
-
-## 核心数据
-
-| 指标 | 值 | 说明 |
-|------|------|------|
-| CCCL 算法总数 | 26 | cub/device/dispatch/tuning/ 下所有 tuning_*.cuh |
-| muh tuning headers | 26 | 1:1 文件对应 ✓ |
-| CCCL 代码行 | 18,094 | 所有 tuning_*.cuh 总和 |
-| muh 代码行 | 3,568 | 19.7% 覆盖率 |
-| CCCL benchmark 注释 | 299 | `ipt_N.tpb_M ... speedup` 格式的数据点 |
-| SM100 模板特化 | 157 | NVIDIA 为 SM100 跑出的最优配置数 |
-| BI-V100 命名 struct | 37 | muh 中 `bi100_*` struct 数量 |
-| 有 bi100 struct 的算法 | 3/26 | reduce(14个), scan(22个), for(1个) |
-| 有 SMEM 保护的算法 | 16/26 | scale_mem_bound 或 while loop |
-
-## 关键发现
-
-### 1. 只有 reduce 和 scan 达到了"READY"状态
-
-reduce 和 scan 是唯一两个同时具备 bi100 命名 struct + SMEM 保护 + 完整 policy_selector 的算法。但即便如此，这些 struct 的值全部是从 SM100 推导的**理论值**，没有一个在 BI-V100 上实测过。
-
-### 2. 其余 24 个算法停留在"inline only"
-
-"inline only" 意味着 muh header 里有 policy_selector，但它的值是硬编码在 if/else 分支里的，不是通过命名 struct 暴露的。gen_patch.py 提取不到这些值（它只认 `struct bi100_*` 模式）。
-
-### 3. CCCL 有 299 个 benchmark 数据点，muh 有 0 个
-
-CCCL 的 benchmark 注释格式完美定义了目标：
-```
-ipt_22.tpb_384.ns_1904.dcid_6.l2w_830.trp_1.ld_0 1.148442 0.997167 1.139902 1.462651
-```
-四个数字 = 四个 problem size 下的加速比。muh 需要在 BI-V100 上产出同样格式的 299 个数据点来填充所有空位。
-
-### 4. 竞赛瓶颈不在代码量而在实测数据
-
-- 代码架构已经搭好（26 个 header + policy_selector + gen_patch 管道）
-- 缺的是 BI-V100 实测数据来替换理论值
-- 没有实测数据，所有 bi100_* struct 的值都是猜的
-
-## 26 算法状态矩阵
-
-| 算法 | CCCL 行 | muh 行 | CCCL BM | SM100 特化 | bi100 struct | SMEM✓ | 状态 |
-|------|---------|--------|---------|-----------|-------------|-------|------|
-| reduce | 478 | 297 | 7 | 6 | 14 | ✓ | ✓ READY |
-| scan | 1,525 | 591 | 18 | 12 | 22 | ✓ | ✓ READY |
-| for | 78 | 51 | 0 | 0 | 1 | ✗ | ⚠ no SMEM |
-| topk | 121 | 113 | 0 | 0 | 0 | ✗ | △ inline |
-| transform | 549 | 185 | 0 | 0 | 0 | ✗ | △ inline |
-| batch_memcpy | 227 | 95 | 0 | 0 | 0 | ✗ | △ inline |
-| select_if | 2,729 | 459 | 84 | 52 | 0 | ✓ | △ inline |
-| radix_sort | 2,381 | 222 | 70 | 0 | 0 | ✓ | △ inline |
-| scan_by_key | 2,008 | 145 | 30 | 17 | 0 | ✓ | △ inline |
-| reduce_by_key | 1,735 | 171 | 32 | 22 | 0 | ✓ | △ inline |
-| unique_by_key | 1,539 | 166 | 29 | 21 | 0 | ✓ | △ inline |
-| three_way_partition | 788 | 99 | 13 | 9 | 0 | ✓ | △ inline |
-| rle_non_trivial_runs | 691 | 68 | 8 | 8 | 0 | ✗ | △ inline |
-| segmented_sort | 640 | 189 | 0 | 0 | 0 | ✓ | △ inline |
-| rle_encode | 626 | 63 | 4 | 7 | 0 | ✗ | △ inline |
-| histogram | 363 | 76 | 4 | 3 | 0 | ✗ | △ inline |
-| segmented_radix_sort | 311 | 48 | 0 | 0 | 0 | ✓ | △ inline |
-| batch_memcpy | 227 | 95 | 0 | 0 | 0 | ✗ | △ inline |
-| merge_sort | 193 | 83 | 0 | 0 | 0 | ✓ | △ inline |
-| segmented_reduce | 189 | 51 | 0 | 0 | 0 | ✗ | △ inline |
-| batched_topk | 186 | 66 | 0 | 0 | 0 | ✓ | △ inline |
-| merge | 180 | 89 | 0 | 0 | 0 | ✓ | △ inline |
-| segmented_scan | 158 | 45 | 0 | 0 | 0 | ✓ | △ inline |
-| adjacent_difference | 118 | 77 | 0 | 0 | 0 | ✓ | △ inline |
-| find_bound_sorted_values | 106 | 47 | 0 | 0 | 0 | ✗ | △ inline |
-| find | 90 | 39 | 0 | 0 | 0 | ✓ | △ inline |
-| transform_tile | 85 | 33 | 0 | 0 | 0 | ✗ | △ inline |
-
-## gen_patch 管道状态
-
-当前 gen_patch.py 跑出来的结果：
-
-```
-READ  reduce:    bi100_plus_float32_o4 → {items:24, threads:512, vec:2}
-READ  scan:      bi100_sm90_float32    → {threads:128, items:24}
-READ  topk:      __inline_topk__       → {threads:512, bits_per_pass:11}
-READ  transform: __inline_transform__  → {bytes_in_flight:64}
-READ  for:       bi100_default         → {threads:256, items:4}
-SKIP  其余 21 个算法: no bi100_* structs
-```
-
-**0 个 patch 生成**——因为 VLLM_INJECTION_POINTS 映射表中的 key 与当前 struct 字段名不匹配。这是管道断裂点。
-
-## CCCL benchmark 源码作为 muh 的输入规范
-
-CCCL bench/reduce/base.cuh 定义了 benchmark 框架：
-- 参数空间：`%RANGE% TUNE_ITEMS_PER_THREAD ipt 7:24:1` / `%RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32`
-- 输出格式：`ipt_N.tpb_M.ipv_K speedup0 speedup1 speedup2 speedup3`
-- 四个 problem size：`Elements{io}` = 2^16, 2^20, 2^24, 2^28
-
-muh 的 bench_bi100.py 已经有 topk 的实测数据（最佳配置：ipt=4, tpb=512, ld=0），
-但 reduce/scan/transform 还没跑。
-
-## CCCL 已有的可直接利用的资产
-
-| 资产类型 | 数量 | 路径 | 用途 |
-|----------|------|------|------|
-| CUB benchmarks | 80 .cu | cccl_upstream/cub/benchmarks/bench/ | 参数空间搜索框架 |
-| CUB tests | 243 .cu | cccl_upstream/cub/test/ | 正确性验证 |
-| CUB examples | 18 .cu | cccl_upstream/cub/examples/ | API 验证 |
-| Thrust examples | 52 .cu | cccl_upstream/thrust/examples/ | 算法验证 |
-| muh schemas | 27 .yaml | muh/schema/ | 参数空间定义 |
-
-总计 420 个 .cu 文件可直接编译运行在 BI-V100 上产出数据。
-
-## 下一步行动
-
-优先级按竞赛权重排序：
-
-1. **reduce 实测** (Output TPS × 16.796 = 83%): 用 bench/reduce/sum.cu 框架，在 BI-V100 上扫描 ipt∈[7,24] × tpb∈{128..1024:32} × ipv∈{1,2,4}
-2. **scan 实测** (decode softmax): 用 bench/scan/exclusive/sum.cu 框架，额外标定 LookbackDelay
-3. **topk 补全** (sampling): 已有部分数据，需要补 batch=4 和 bits_per_pass 对比
-4. **gen_patch 闭环**: 修复 VLLM_INJECTION_POINTS 映射，让 gen_patch 真正产出可用 patch
-5. **50+ 功能测试**: 在 patch 后的 vllm 上跑竞赛功能验证

CCCL_MUH_PARITY_AUDIT.md

@@ -1,134 +0,0 @@
-================================================================================
-CCCL vs muh 精确比对审计报告
-================================================================================
-
-### 1. scale_mem_bound 函数 parity check
-------------------------------------------------------------
-  float32 (CCCL SM100 reduce)         CCCL=( 16i, 512t,tile= 32768B) muh=( 16i, 512t,tile= 32768B) ✓
-  float64 (CCCL SM100 reduce)         CCCL=(  8i, 640t,tile= 40960B) muh=(  8i, 640t,tile= 40960B) ✓
-  accum8 (CCCL SM100 reduce)          CCCL=(  7i, 512t,tile= 28672B) muh=(  7i, 512t,tile= 28672B) ✓
-  scan 4B (CCCL SM100 scan)           CCCL=( 22i, 384t,tile= 33792B) muh=( 22i, 384t,tile= 33792B) ✓
-  scan 8B (CCCL SM100 scan)           CCCL=( 11i, 416t,tile= 36608B) muh=( 11i, 416t,tile= 36608B) ✓
-  det float32 SM90                    CCCL=( 13i, 224t,tile= 11648B) muh=( 13i, 224t,tile= 11648B) ✓
-  det float64 SM86                    CCCL=(  5i, 128t,tile=  5120B) muh=(  5i, 128t,tile=  5120B) ✓
-  1-byte type                         CCCL=( 32i, 256t,tile=  8192B) muh=( 32i, 256t,tile=  8192B) ✓
-  2-byte type                         CCCL=( 32i, 256t,tile= 16384B) muh=( 32i, 256t,tile= 16384B) ✓
-  16-byte type (int128)               CCCL=(  4i, 256t,tile= 16384B) muh=(  4i, 256t,tile= 16384B) ✓
-  SMEM cap test (should trigger)      CCCL=(  8i, 768t,tile= 49152B) muh=(  8i, 768t,tile= 49152B) ✓
-  → scale_mem_bound: FULL PARITY ✓
-
-### 2. reduce tuning: CCCL SM100值 → BI-V100 scale_mem_bound适配后
-------------------------------------------------------------
-  CCCL benchmarked on SM100 → muh should use scale_mem_bound for BI-V100
-  Key: reduce loads to REGISTERS not SMEM → SMEM cap rarely triggers
-
-  float32_plus_o4           @4B: scaled=(16i,  512t) tile= 32768B (66.7%)
-  float32_plus_o4           @8B: scaled=( 8i,  512t) tile= 32768B (66.7%)
-  float64_plus_o4           @4B: scaled=(16i,  640t) tile= 40960B (83.3%)
-  float64_plus_o4           @8B: scaled=( 8i,  640t) tile= 40960B (83.3%)
-  accum8_plus_o4            @4B: scaled=(15i,  512t) tile= 30720B (62.5%)
-  accum8_plus_o4            @8B: scaled=( 7i,  512t) tile= 28672B (58.3%)
-  accum8_plus_o8            @4B: scaled=(15i,  512t) tile= 30720B (62.5%)
-  accum8_plus_o8            @8B: scaled=( 7i,  512t) tile= 28672B (58.3%)
-  det_float32_sm90          @4B: scaled=(13i,  224t) tile= 11648B (23.7%)
-  det_float32_sm90          @8B: scaled=( 6i,  224t) tile= 10752B (21.9%)
-  det_float32_sm86          @4B: scaled=( 6i,  224t) tile=  5376B (10.9%)
-  det_float32_sm86          @8B: scaled=( 3i,  224t) tile=  5376B (10.9%)
-  det_float64_sm86          @4B: scaled=(11i,  128t) tile=  5632B (11.5%)
-  det_float64_sm86          @8B: scaled=( 5i,  128t) tile=  5120B (10.4%)
-  default_fallback          @4B: scaled=(16i,  256t) tile= 16384B (33.3%)
-  default_fallback          @8B: scaled=( 8i,  256t) tile= 16384B (33.3%)
-
-### 3. muh bi100 reduce当前值 vs CCCL参考
-------------------------------------------------------------
-  muh改用了更大的items (24 vs SM100的16)来补偿16 SMs
-  这是对的——reduce加载到寄存器,SMEM不是瓶颈
-
-  ★ float32 plus (paged_attention score reduction — 83% weight):
-    CCCL SM100:  items=16, threads=512, vec=2
-    muh BI-V100: items=24, threads=512, vec=2
-    理由: 16 SMs vs 148 SMs, 每个CTA需要处理更多数据
-    tile对比: SM100=512*16*4=32768B | BI-V100=512*24*4=49152B (exactly 48KB)
-    → items=24 用满了SMEM → 合理但有风险,如果BlockReduce实际占SMEM则溢出
-    → 但注释说reduce不用BlockLoad(loads to registers) → 安全
-
-### 4. scan tuning: CCCL SM100 → BI-V100 SMEM约束
-------------------------------------------------------------
-  Scan DOES use BlockLoad staging in SMEM → tile_bytes ≤ 49152 is HARD
-
-  lookback_1B_o4       @1B: tpb= 512 ipt=18 tile=  9216B ✓
-  lookback_1B_o4       @2B: tpb= 512 ipt=18 tile= 18432B ✓
-  lookback_1B_o4       @4B: tpb= 512 ipt=18 tile= 36864B ✓
-  lookback_1B_o4       @8B: tpb= 512 ipt=18 tile= 73728B ✗ OVERFLOW → max_items=12
-  lookback_2B_o4       @1B: tpb= 512 ipt=13 tile=  6656B ✓
-  lookback_2B_o4       @2B: tpb= 512 ipt=13 tile= 13312B ✓
-  lookback_2B_o4       @4B: tpb= 512 ipt=13 tile= 26624B ✓
-  lookback_2B_o4       @8B: tpb= 512 ipt=13 tile= 53248B ✗ OVERFLOW → max_items=12
-  lookback_4B_o4       @1B: tpb= 384 ipt=22 tile=  8448B ✓
-  lookback_4B_o4       @2B: tpb= 384 ipt=22 tile= 16896B ✓
-  lookback_4B_o4       @4B: tpb= 384 ipt=22 tile= 33792B ✓
-  lookback_4B_o4       @8B: tpb= 384 ipt=22 tile= 67584B ✗ OVERFLOW → max_items=16
-  lookback_8B_o4       @1B: tpb= 416 ipt=23 tile=  9568B ✓
-  lookback_8B_o4       @2B: tpb= 416 ipt=23 tile= 19136B ✓
-  lookback_8B_o4       @4B: tpb= 416 ipt=23 tile= 38272B ✓
-  lookback_8B_o4       @8B: tpb= 416 ipt=23 tile= 76544B ✗ OVERFLOW → max_items=14
-  lookback_1B_o8       @1B: tpb= 384 ipt=14 tile=  5376B ✓
-  lookback_1B_o8       @2B: tpb= 384 ipt=14 tile= 10752B ✓
-  lookback_1B_o8       @4B: tpb= 384 ipt=14 tile= 21504B ✓
-  lookback_1B_o8       @8B: tpb= 384 ipt=14 tile= 43008B ✓
-  lookback_4B_o8       @1B: tpb= 416 ipt=19 tile=  7904B ✓
-  lookback_4B_o8       @2B: tpb= 416 ipt=19 tile= 15808B ✓
-  lookback_4B_o8       @4B: tpb= 416 ipt=19 tile= 31616B ✓
-  lookback_4B_o8       @8B: tpb= 416 ipt=19 tile= 63232B ✗ OVERFLOW → max_items=14
-  lookback_8B_o8       @1B: tpb= 320 ipt=22 tile=  7040B ✓
-  lookback_8B_o8       @2B: tpb= 320 ipt=22 tile= 14080B ✓
-  lookback_8B_o8       @4B: tpb= 320 ipt=22 tile= 28160B ✓
-  lookback_8B_o8       @8B: tpb= 320 ipt=22 tile= 56320B ✗ OVERFLOW → max_items=19
-
-  关键发现:
-  - scan lookback_4B_o4: items=22, threads=384 → tile@4B=33792 ✓ tile@8B=67584 ✗
-  - scan lookback_8B_o4: items=23, threads=416 → tile@8B=76544 ✗
-  - 这些值在SM100上是安全的(228KB SMEM),但在BI-V100(48KB)上必须降级
-  - muh已经做了降级(用scale_mem_bound),但需要验证降级后的值是否正确
-
-### 5. CCCL benchmark format解析
-------------------------------------------------------------
-  NVIDIA的benchmark注释格式:
-  ipt_<items>.tpb_<threads>.ns_<delay>.dcid_<algo>.l2w_<latency>.trp_<transpose>.ld_<load>
-  后跟4个浮点数: 在[2^16, 2^20, 2^24, 2^28]四个problem size下的speedup
-
-  dcid映射:
-    0 = no_delay
-    1 = fixed_delay
-    2 = exp_backoff
-    3 = exp_backoff_jitter
-    4 = exp_backoff_jitter_window
-    5 = exp_backon_jitter_window
-    6 = exp_backon_jitter
-    7 = exp_backon
-
-### 6. 竞赛关键路径优先级
-------------------------------------------------------------
-  Token吞吐加权值 = Output_TPS × 16.796 + Input_TPS × 2.799 + Cache_TPS × 0.56
-  → Output_TPS权重83%, Input_TPS权重14%, Cache_TPS权重3%
-
-  decode热路径 (Output TPS):
-    1. paged_attention score reduction → reduce (DONE: muh tuned)
-    2. softmax denominator prefix-sum → scan (DONE: muh tuned)
-    3. top-k/top-p sampling → topk/radix_sort (DONE: muh tuned)
-    4. RMSNorm/SiLU/RoPE element-wise → transform (DONE: muh tuned)
-
-  prefill热路径 (Input TPS):
-    5. flash_attention → scan + reduce
-    6. MoE expert routing → select_if + reduce_by_key
-
-  cache热路径 (Cache TPS):
-    7. KV cache block copy → batch_memcpy (DONE: muh tuned)
-
-### 7. 待验证的关键问题
-------------------------------------------------------------
-  1. reduce items=24: 虽然loads to registers, 但实际BlockReduce<WARP_REDUCTIONS>的SMEM用量需要确认
-  2. scan delay参数: 0.5x/0.6x缩放是启发式, 需要BI-V100实测L2 write latency
-  3. LOAD_LDG vs LOAD_DEFAULT: topk bench显示BI-V100上LOAD_DEFAULT更快, reduce/scan可能同理
-  4. SM count=16 → wave efficiency: 所有tuning都需要重新算occupancy
-  5. transform bytes_in_flight: 从18GB/s改为56GB/s后items需要相应增大

CCCL_PATTERN_MAP.md

@@ -1,104 +0,0 @@
-# CCCL → vllm Kernel Pattern Mapping
-## BI-V100 Competition Reference
-
-### Pattern 1: Multi-field Reduction (paged_attention)
-
-**CCCL source**: `thrust/examples/bounding_box.cu`, `summary_statistics.cu`
-**vllm kernel**: `paged_attn.py` → ixformer paged_attention_v1/v2
-
-```
-CCCL:  transform_reduce(begin, end, unary_op, init, binary_op)
-vllm:  for each KV block: score = Q·K, max_score = reduce_max, exp_sum = reduce_sum
-```
-
-**Tuning surface**:
-- `_PARTITION_SIZE`: controls how many KV tokens per CTA in V2 mode
-- V1/V2 dispatch threshold: `total_tiles vs 2 × sm_count`
-- BI-V100: 16 SMs → V2 beneficial when seq_len > 1024 (2 waves of 16 CTAs × 512 partition)
-
-**CCCL parameter**: `ReducePassPolicy{threads=512, items=24, vec=2, WARP_REDUCTIONS, LDG}`
-
-### Pattern 2: Prefix Scan + Transform (softmax)
-
-**CCCL source**: `thrust/examples/simple_moving_average.cu`, `cub/benchmarks/bench/scan/exclusive/sum.cu`
-**vllm kernel**: `prefix_prefill.py` context_attention_fwd_kernel
-
-```
-CCCL:  inclusive_scan(begin, end, output, plus<float>)
-vllm:  for each BLOCK_N chunk: qk = Q·K, m_new = max(m_old, max(qk)),
-       l_new = l_old * exp(m_old - m_new) + sum(exp(qk - m_new))
-```
-
-**Tuning surface**:
-- `BLOCK_M`: Q tile rows (32 or 64 for BI-V100)
-- `BLOCK_N`: K/V sweep width (32 or 64)
-- `NUM_WARPS`: 4 (16 SMs don't benefit from 8 warps per CTA)
-- `num_stages`: 1 (no cp.async) or 2 (software pipeline)
-
-**CCCL parameter**: `ScanLookbackPolicy{threads=384, items=22, WARP_TRANSPOSE, DEFAULT, WARP_SCANS, {backon_jitter_window, 952, 415}}`
-
-### Pattern 3: Transform (activation functions)
-
-**CCCL source**: `cub/benchmarks/bench/transform/babelstream.cu`
-**vllm kernel**: Triton SiLU, GeLU, RMSNorm kernels (via `_custom_ops.py`)
-
-```
-CCCL:  transform(begin, end, output, silu_op)    // x * sigmoid(x)
-vllm:  @triton.jit def silu_kernel(x): tl.sigmoid(x) * x
-```
-
-**Tuning surface**:
-- `bytes_in_flight`: 64KB on BI-V100 (56 GB/s per-SM × 1100ns latency)
-- Triton `num_stages=2` maps to BIF=64KB (2× prefetch window)
-- `SMEM = 49152` (fixed by _custom_ops.py)
-
-**CCCL parameter**: `TransformPrefetchPolicy{threads=256, bif=64KB, prefetch_stride=128}`
-
-### Pattern 4: TopK (sampling)
-
-**CCCL source**: `cub/benchmarks/bench/topk/keys.cu`
-**vllm kernel**: sampling_kernels (precompiled .so)
-
-```
-CCCL:  DeviceTopk::TopK(keys, k, output)
-vllm:  ixformer topk_sampling → radix_sort + select partial
-```
-
-**Tuning surface** (via .so, limited):
-- `bits_per_pass`: 11 for float32 (32 bits / 3 passes)
-- Thread count: 512 (baked into .so)
-
-### Pattern 5: Triton Flash Attention (all patterns combined)
-
-**CCCL source**: All of the above + `cub/agent/agent_scan.cuh` union SMEM model
-**vllm kernel**: `triton_flash_attention.py`
-
-```
-Q_resident × K_streaming × softmax_online → Output
-= transform_reduce (Q·K) + scan (softmax) + transform (V matmul)
-```
-
-**Tuning surface**: 17 existing + 19 new autotune configs from gen_config.py
-**Key configs for BI-V100**:
-```python
-# Best for long context (seq_len > 4096):
-Config(BLOCK_M=64, BLOCK_N=64, num_warps=4, num_stages=2)  # 40KB SMEM, 1 CTA/SM
-
-# Best for short context (seq_len < 1024):
-Config(BLOCK_M=32, BLOCK_N=32, num_warps=2, num_stages=2)  # 32KB SMEM, 2 CTAs/SM
-```
-
----
-
-### CCCL Asset Utilization Summary
-
-| CCCL Asset | Files | Used for BI-V100 | Competition Impact |
-|-----------|-------|-------------------|-------------------|
-| Tuning headers (26) | 18094 lines | 3568 lines (20%) | P0: reduce/scan/transform |
-| CUB benchmarks (80) | reduce/scan/topk/transform | benchmark framework | P0: parameter search |
-| Thrust examples (52) | summary_stats/bounding_box/norm | pattern mapping | P1: architecture understanding |
-| CUB tests (243) | correctness verification | 0% (need BI-V100) | P2: correctness |
-| libcudacxx (1463) | type traits, atomics | implicit (via CUB) | Infra |
-
-**Total usable CCCL assets**: 5205 files in cccl_upstream
-**Competition-critical subset**: ~30 files (5 tuning headers + 10 benchmarks + 15 examples)

CCCL_TUNING_GAP_REPORT.md

@@ -1,89 +0,0 @@
-# CCCL ↔ muh Tuning Header Gap Report
-
-> **Generated**: 2026-08-06 (auto-analyzed from source code)
-> **Source of truth**: `cccl_upstream/cub/cub/device/dispatch/tuning/tuning_*.cuh`
-> **muh headers**: `muh/include/muh/tuning/tuning_*.cuh`
-
-## Executive summary
-
-- **26 algorithms** have both CCCL original and muh BI-V100 tuning headers.
-- muh covers **19% of CCCL lines** (3568 / 18094).
-- CCCL contains **294 benchmark annotations** across all algorithms. muh has **1 benchmarked algorithm** (scan, partial).
-- The **#1 gap** is not code coverage — it's the absence of BI-V100 benchmark data in `ipt_N.tpb_M speedup` format.
-
-## Per-algorithm coverage
-
-| Algorithm | CCCL lines | muh lines | Coverage | CCCL bench pts | muh bi100 structs | muh benchmarked? |
-|-----------|-----------|----------|----------|---------------|-------------------|-----------------|
-| reduce | 478 | 297 | 62% | 6 | 14 | ✗ |
-| scan | 1525 | 591 | 38% | 16 | 22 | ✓ (partial) |
-| topk | 121 | 113 | 93% | 0 | 0 | ✗ |
-| radix_sort | 2381 | 222 | 9% | 70 | 0 | ✗ |
-| select_if | 2729 | 459 | 16% | 82 | 0 | ✗ |
-| scan_by_key | 2008 | 145 | 7% | 30 | 0 | ✗ |
-| reduce_by_key | 1735 | 171 | 9% | 32 | 0 | ✗ |
-| unique_by_key | 1539 | 166 | 10% | 29 | 0 | ✗ |
-| three_way_partition | 788 | 99 | 12% | 13 | 0 | ✗ |
-| rle_non_trivial_runs | 691 | 68 | 9% | 8 | 0 | ✗ |
-| segmented_sort | 640 | 189 | 29% | 0 | 0 | ✗ |
-| rle_encode | 626 | 63 | 10% | 4 | 0 | ✗ |
-| transform | 549 | 185 | 33% | 0 | 0 | ✗ |
-| histogram | 363 | 76 | 20% | 4 | 0 | ✗ |
-| segmented_radix_sort | 311 | 48 | 15% | 0 | 0 | ✗ |
-| batch_memcpy | 227 | 95 | 41% | 0 | 0 | ✗ |
-| batched_topk | 186 | 66 | 35% | 0 | 0 | ✗ |
-| merge_sort | 193 | 83 | 43% | 0 | 0 | ✗ |
-| merge | 180 | 89 | 49% | 0 | 0 | ✗ |
-| segmented_reduce | 189 | 51 | 26% | 0 | 0 | ✗ |
-| segmented_scan | 158 | 45 | 28% | 0 | 0 | ✗ |
-| adjacent_difference | 118 | 77 | 65% | 0 | 0 | ✗ |
-| find | 90 | 39 | 43% | 0 | 0 | ✗ |
-| find_bound_sorted_values | 106 | 47 | 44% | 0 | 0 | ✗ |
-| transform_tile | 85 | 33 | 38% | 0 | 0 | ✗ |
-| for | 78 | 51 | 65% | 0 | 1 | ✗ |
-| **TOTAL** | **18094** | **3568** | **19%** | **294** | **37** | **1/26** |
-
-## Reduce: CCCL SM100 → muh BI-V100 divergence analysis
-
-### SM100 benchmark annotations in CCCL
-```
-ipt_15.tpb_512.ipv_2  1.020  1.000  1.018  1.058  (geo=1.024) — accum8, offset4
-ipt_15.tpb_512.ipv_1  1.019  1.000  1.017  1.057  (geo=1.023) — accum8, offset8
-ipt_16.tpb_512.ipv_2  1.061  1.000  1.065  1.167  (geo=1.072) — float32, offset4
-ipt_16.tpb_640.ipv_1  1.018  1.000  1.016  1.057  (geo=1.022) — float64, offset4
-ipt_13.tpb_224        1.107  1.010  1.097  1.317  (geo=1.127) — deterministic float32 (sm90)
-ipt_6.tpb_224         1.034  1.000  1.032  1.091  (geo=1.039) — deterministic float32 (sm86)
-```
-
-### Key divergences
-
-| Parameter | CCCL SM100 | muh BI-V100 | Rationale | Risk |
-|-----------|-----------|------------|-----------|------|
-| float32+plus items | 16 | 24 | Compensate for 16 vs 148 SMs | Unvalidated: may hurt L1 hit rate |
-| float64+plus threads | 640 | 384 | Clean 12-warp config | May underutilize vs 20-warp original |
-| float64+plus vec | 1 | 2 | 16B vectorized loads | Alignment risk with non-contiguous data |
-| det float32 items | 13 | 32 | More work per CTA on 16 SMs | 2.5× register pressure increase |
-| accum1/2/16 | absent | added | Extrapolated from scaling | Not in CCCL SM100, completely theoretical |
-
-## Scan: lookback delay calibration gap
-
-CCCL SM100 lookback delay parameters (from benchmark annotations):
-- `delay_ns` range: 228 – 1904 ns
-- `dcid` (delay constructor ID) range: 1 – 7
-- `l2_write_latency` range: 520 – 965 ns
-
-These are calibrated on SM100's 50MB L2 cache. BI-V100 has 6MB L2 → delay parameters need re-calibration. Current muh values use heuristic scaling (SM100 × 0.5 for ns, × 0.6 for l2w) without hardware validation.
-
-## Priority action items (by Output TPS impact)
-
-| # | Algorithm | CCCL bench pts needed | vllm hot path | Weight |
-|---|-----------|----------------------|---------------|--------|
-| 1 | reduce | 6 | paged_attention score reduction | 83% |
-| 2 | scan | 16 (8 remaining) | softmax denominator | 83% |
-| 3 | topk | 0 (format from radix_sort) | vocab=152064 sampling | 83% |
-| 4 | radix_sort | 70 | logit sorting for top-k/top-p | 83% |
-| 5 | select_if | 82 | top-p token filtering | 83% |
-| 6 | transform | 0 (no CCCL benches) | RMSNorm/SiLU/RoPE | 10-15% |
-| 7 | scan_by_key | 30 | per-sequence softmax | ~5% |
-| 8 | reduce_by_key | 32 | per-sequence aggregation | ~3% |
-| 9 | batch_memcpy | 0 | KV cache block copy | 3% |

CODEPATH_MAP.md

@@ -1,193 +0,0 @@
-# 代码路径时序图 — 从HTTP请求到GPU kernel的完整链路
-
-## 一、请求入口到引擎调用
-
-```
-HTTP POST /v1/chat/completions
-  │
-  ├─ api_server.py → FastAPI route handler
-  │    └─ serving_chat.py:create_chat_completion() [line ~140]
-  │         ├─ protocol.py:ChatCompletionRequest.model_validate()
-  │         │    └─ max_completion_tokens → max_tokens 映射 [line 418]
-  │         │    └─ extra="allow" (Sub168用extra="forbid"导致400)
-  │         │
-  │         ├─ chat_utils.py → 消息格式化 + 多模态处理
-  │         │    └─ content=None容错 (Sub168这里崩)
-  │         │
-  │         ├─ serving_chat.py [line 175-213] → enable_thinking逻辑
-  │         │    ├─ tool_choice=auto + tools存在 → enable_thinking=False
-  │         │    ├─ thinking.type=disabled → enable_thinking=False
-  │         │    └─ 默认 → enable_thinking=True
-  │         │
-  │         ├─ serving_chat.py [line 250-252] → n值检查
-  │         │    └─ n>2 → 400 (n=2允许传入引擎)
-  │         │
-  │         └─ engine_client.generate() [line 355]
-  │              └─ try/except ValueError + catch-all Exception
-  │
-  ├─ computility-run.yaml → vLLM启动参数
-  │    ├─ --max-num-seqs 2 (防止n=2崩溃)
-  │    ├─ --max-model-len 80000
-  │    ├─ --enforce-eager (禁用CUDA Graph)
-  │    ├─ --enable-prefix-caching
-  │    └─ --tool-call-parser qwen3_coder
-  │
-  └─ 如果引擎crash → 后续所有请求Connection Refused
-       (Sub508的根因: t2_n_2触发, 30个FAIL级联)
-```
-
-## 二、模型前向传播 — 逐层链路
-
-```
-Qwen3_5ForCausalLM.forward()  [qwen3_5.py line 1214]
-  │
-  └─ Qwen3_5Model.forward()  [line 1094]
-       │
-       ├─ embed_tokens(input_ids)
-       │
-       └─ for layer in self.layers:  # 36层 (Qwen3.6-27B典型配置)
-            │
-            ├─ GemmaRMSNorm(hidden_states, residual)
-            │    └─ ☆ 可用ixformer: fused_add_rms_norm(input, residual, weight, eps)
-            │
-            ├─ [linear_attention层] GatedDeltaNet.forward() [line 407]
-            │    │
-            │    ├─ CoreX dispatch尝试 [line 416-425]
-            │    │    └─ _use_corex_gdn=False (base image无corex_gdn模块)
-            │    │
-            │    └─ _pytorch_forward() [line 435]  ← 当前执行路径
-            │         │
-            │         ├─ 投影: in_proj_qkv, in_proj_z, in_proj_b, in_proj_a
-            │         │    └─ ☆ 每个是F.linear → 可用ixformer.matmul
-            │         │
-            │         ├─ [prefill] 逐序列循环 [line 463-555]
-            │         │    │
-            │         │    ├─ F.conv1d (causal conv)
-            │         │    │    └─ ☆ 可用ixformer.conv2d (需reshape)
-            │         │    │
-            │         │    ├─ F.silu → ☆ 可用ixformer.silu_and_mul
-            │         │    │
-            │         │    ├─ g计算: -A_log.exp() * softplus(a+dt_bias)
-            │         │    │    └─ 当前: clamp(-8,4)后exp, softplus.clamp(max=10)
-            │         │    │
-            │         │    └─ _torch_chunk_gated_delta_rule() [line 152-247]
-            │         │         │
-            │         │         ├─ g.clamp(-5,2).cumsum(-1).clamp(-20,20) ← NaN修复点
-            │         │         ├─ decay_mask = exp(g差) ← 所有exp在clamp后
-            │         │         ├─ attn矩阵: k_beta @ key.T * decay_mask
-            │         │         │    └─ ☆ 三角求解循环 → 无法用ixformer加速
-            │         │         │       (这是纯序列依赖: attn[i] += attn[i,:i] @ attn[:i,:i])
-            │         │         ├─ state更新循环: for i in chunks [line 219-232]
-            │         │         │    ├─ q @ k.T * decay ← ☆ ixformer.matmul可加速
-            │         │         │    ├─ q * exp(g) @ state ← ☆ ixformer.matmul可加速
-            │         │         │    └─ state更新: state * exp(g) + k.T @ v_new
-            │         │         │         └─ ☆ ixformer.matmul可加速
-            │         │         └─ 最终: core_out → transpose → to(dtype)
-            │         │
-            │         ├─ [decode] 单token路径 [line 558-638]
-            │         │    ├─ _torch_causal_conv1d_update
-            │         │    │    └─ 逐通道点积 → ☆ ixformer.gemv可加速
-            │         │    ├─ g_t = g.clamp(-20,2).exp_() ← NaN修复点
-            │         │    ├─ temporal_state.mul_(g_t) ← 状态衰减
-            │         │    ├─ torch.bmm(k, state) ← ☆ ixformer.matmul可加速
-            │         │    └─ state.baddbmm_(k, delta) ← ☆ ixformer.matmul可加速
-            │         │
-            │         └─ GemmaRMSNorm + out_proj
-            │              └─ ☆ ixformer.rms_norm + ixformer.matmul
-            │
-            ├─ [full_attention层] Qwen3_5FullAttention.forward() [line 737]
-            │    └─ 标准vLLM Attention → XFormers后端
-            │         └─ ☆ 已使用ixformer.flash_attn_func (base image配置)
-            │
-            ├─ GemmaRMSNorm(hidden_states, residual)
-            │    └─ ☆ ixformer.fused_add_rms_norm
-            │
-            └─ [MLP/MoE] Qwen3_5MLP 或 Qwen3_5MoeSparseBlock
-                 │
-                 ├─ [MLP] gate_up_proj → silu_and_mul → down_proj
-                 │    └─ ☆ 全部可用ixformer: matmul + silu_and_mul + matmul
-                 │
-                 └─ [MoE] Qwen3_5MoeSparseBlock.forward() [line 974]
-                      ├─ gate(hidden) → router_logits
-                      ├─ softmax → topk → renormalize (纯PyTorch, 无硬件加速)
-                      ├─ _pure_pytorch_experts() [line 897]
-                      │    ├─ [decode T=1] 批量GEMM: 3次kernel launch
-                      │    │    └─ F.linear(x, w13_sel.reshape(-1,H)) ← ☆ ixformer.matmul
-                      │    │    └─ F.silu(gate) * up ← ☆ ixformer.silu_and_mul (需reshape)
-                      │    │    └─ torch.bmm(w2_sel, act) ← ☆ ixformer.matmul
-                      │    └─ [prefill] 逐expert循环 ← 性能瓶颈
-                      │         └─ 每个expert: F.linear × 2 + silu
-                      │              └─ ☆ 可用ixformer.matmul但循环开销不变
-                      └─ shared_expert: gate_up → silu_and_mul → down → sigmoid gate
-                           └─ ☆ 全部可用ixformer
-```
-
-## 三、ixformer可用原语 vs 当前使用情况
-
-| ixformer原语 | 签名 | 当前是否使用 | 可替换的PyTorch调用 |
-|-------------|------|------------|-------------------|
-| `matmul` | `matmul(input, other, out, transa, transb, alpha, beta)` | ❌ 未使用 | F.linear, torch.mm, torch.bmm, @ |
-| `softmax` | `softmax(input, dim)` | ❌ 未使用 | torch.softmax (MoE路由) |
-| `rms_norm` | `rms_norm(input, weight, output, eps)` | ❌ 未使用 | GemmaRMSNorm内部 |
-| `fused_add_rms_norm` | `fused_add_rms_norm(input, residual, weight, eps, scale)` | ❌ 未使用 | residual + layernorm 两步 |
-| `silu_and_mul` | `silu_and_mul(input, output)` | ❌ 未使用 | SiluAndMul层, F.silu(g)*up |
-| `conv2d` | `conv2d(input, weight, bias, stride, padding, dilation, groups)` | ❌ 未使用 | F.conv1d (causal conv) |
-| `flash_attn_func` | `flash_attn_func(q, k, v, dropout_p, softmax_scale, causal)` | ✅ XFormers后端使用 | full_attention层 |
-| `gemv` | `gemv(x, A)` | ❌ 未使用 | decode路径小矩阵乘 |
-| `scaled_dot_product_attention` | `sdpa(query, key, value, attn_mask, dropout_p, is_causal)` | ❌ 未使用 | 可替代chunk内QK^T计算 |
-
-**关键发现：9个可用原语中只有1个（flash_attn_func）被使用，而且不是我们的代码使用的——是base image的XFormers后端自动调用的。我们的代码对ixformer的利用率是0%。**
-
-## 四、Sub168 vs Sub508 性能差距的代码解释
-
-```
-Sub168 (8.49s for d01):
-  base image native qwen3_5.py
-  ├─ corex_gdn: 使用libcorex_gdn.so的fused GDN kernel ← 不存在于我们的base image
-  ├─ corex_moe: 使用libcorex_moe.so的fused MoE kernel ← 不存在于我们的base image
-  └─ 所有底层ops由ixformer后端加速 (matmul/rms_norm/softmax等)
-
-Sub508 (95.85s for d01):
-  我们的自定义 qwen3_5.py
-  ├─ GatedDeltaNet: 纯PyTorch (cumsum→exp→NaN→nan_to_num→全零)
-  ├─ MoE: 纯PyTorch循环 (每expert单独F.linear)
-  └─ 底层ops全部用PyTorch默认kernel (未调用ixformer)
-```
-
-## 五、优化路径 — 用ixformer原语替换PyTorch
-
-### 立即可做 (不改算法, 只换kernel):
-1. **matmul**: 所有F.linear/torch.bmm/@ → ixformer.matmul
-2. **silu_and_mul**: MLP和MoE的silu*gate → ixformer.silu_and_mul
-3. **rms_norm**: GemmaRMSNorm内部 → ixformer.rms_norm
-4. **fused_add_rms_norm**: residual+norm两步 → 一步fused
-5. **softmax**: MoE路由softmax → ixformer.softmax
-
-## 六、功能测试FAIL根因分析（6个非crash FAIL）
-
-```
-FAIL类型A: NaN导致模型输出质量问题 (修NaN后自愈)
-├─ d03_tool_call: tools=0 — 模型不能输出<tool_call> XML
-├─ d07_reasoning_plus_content: content[0] — 模型不输出</think>
-├─ d10_thinking_disable_ctk: 乱码 — 模型logits被NaN扭曲
-├─ t1a_thinking_true: reasoning[0] — output.text为空→parser返回空
-└─ t1c_thinking_default: reasoning[0] — 同上
-
-FAIL类型B: 请求处理层问题
-└─ d05_multimodal: HTTP 400 — 多模态请求验证失败
-
-FAIL类型C: 引擎crash级联 (修max-num-seqs=2后自愈)
-└─ t2_n_2 → t3/t4/t5/t6/t7/t8/t9/t10/t12/t13/t14/t15/t16 全部HTTP 500 (25个)
-
-当前代码状态:
-  NaN修复: ✅ cumsum前clamp[-5,2] + 后clamp[-20,20] + A_log clamp[-8,4]
-  引擎防崩: ✅ max-num-seqs=2 + catch-all Exception
-  ixformer加速: ✅ matmul/bmm/softmax接入12处热路径
-  reasoning parser: ✅ qwen3已注册，部署正确
-  tool parser: ✅ qwen3_coder已注册，adjust_request禁thinking
-  
-预期: NaN修复后模型质量恢复 → 类型A的5个FAIL自愈
-      max-num-seqs=2 → 类型C的25个FAIL自愈
-      剩余: d05_multimodal需要单独debug
-      预估: 45/51 PASS (88%)
-```

COMP168_DIAGNOSIS.md

@@ -1,165 +0,0 @@
-# comp 168 Docker 诊断 → .so 开发清单
-
-> 基于 `2d5232c5d6bc` (comp 168 docker log, 3786 行)
-> 当前 HEAD: `b25fc53e` (414 commits)
-
-## 一、comp 168 日志三大致命问题
-
-| # | 错误 | 出现次数 | 根因 | 状态 |
-|---|------|----------|------|------|
-| 1 | `GDN NaN frac=0.9998` | 16次(layer 0-4) | 我们的 GDN prefill 实现产生 NaN → replace with zeros → 模型质量归零 | **P0 未修** |
-| 2 | `vllm_moe_topk_softmax not found` | 39次 | `ixformer.functions` 没有 Python binding → fallback to Python for 循环 | **P0 需 .so** |
-| 3 | `CUDA OOM 32 MiB` | 17次 | `max_model_len=100000` 超过 KV cache 容量 → engine 死亡 | ✅ 已修为 80000 |
-
-## 二、真机探测确认的事实
-
-从你贴的真机 probe 输出：
-
-```
-ixformer.functions 有:
-  ✓ silu_and_mul, rms_norm, fused_add_rms_norm, rotary_embedding
-  ✓ flash_attn_*, vllm_single_query_cached_kv_attention_v2
-  ✓ vllm_cache_ops_reshape_and_cache, vllm_swap_blocks, vllm_copy_cache
-  ✗ vllm_moe_topk_softmax (不存在!)
-  ✗ moe_compute_token_index_api (不存在!)
-  ✗ moe_w16a16_group_gemm (不存在!)
-
-libixformer.so 中:
-  ✓ 上述函数全部存在 (C++ 符号, xllm 的 ixformer.h 声明了它们)
-  但 Python binding (_C.so) 没有暴露
-```
-
-**结论**: MoE 7 步 pipeline 中的 topk_softmax / gen_idx / expand / group_gemm / combine 全部需要通过 `ix_moe_bridge.so` 桥接。
-
-## 三、需要开发/修复的 .so 清单
-
-### SO-1: `ix_moe_bridge.so` (MoE 7步 pipeline) — ✅ 代码已有，需真机编译验证
-
-**源码**: `ex_engine/csrc/ix_moe_bridge.cpp` (258行)
-**编译**: `ex_engine/precompile_ix_bridge.py` → `torch.utils.cpp_extension.load(-lixformer)`
-**状态**: 代码写好了，Dockerfile 有 build step，但从未在真机验证过编译成功
-
-真机验证命令:
-```bash
-cd /workspace/ex_engine
-python3 precompile_ix_bridge.py
-ls -la build/ix_moe_bridge*.so
-python3 -c "import torch; from torch.utils.cpp_extension import load; m=load('test', sources=['csrc/ix_moe_bridge.cpp'], extra_ldflags=['-L/usr/local/corex/lib64/python3/dist-packages/ixformer', '-lixformer']); print(dir(m))"
-```
-
-### SO-2: GDN prefill 修复 — **P0 最高优先级**
-
-**现状**: 我们的 `_torch_chunk_gated_delta_rule` 在 fp16 下产生 99.98% NaN
-**参考**: `upstream_ref/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp` (576行)
-
-关键差异:
-- xllm 用 `fp32` accumulation: `decay_mask = ... .exp().float()` 
-- xllm 用 `torch::matmul` 而不是自定义 chunk kernel
-- xllm 的 recurrent state 管理有精确的 `clamp(-20, 20)` 限制
-
-**解决方案**: 不写新 .so，而是从 xllm 搬运 GDN 的 PyTorch 实现（C++ torch ops, 全 fp32 accumulation），替换我们的 chunk kernel。
-
-### SO-3: `_custom_ops.py` patch — ✅ 已有 fallback 逻辑
-
-base image 的 `_custom_ops.py` 调用 `ixf_F.vllm_moe_topk_softmax` 时会报错。
-但 comp 168 的 base 镜像绕过了 `_custom_ops`，直接走 `corex_moe.py` 的 7 步 pipeline。
-
-**如果 base 有 corex_moe.py**: 不需要 patch
-**如果 base 没有 corex_moe.py**: 我们的版本 + ix_moe_bridge.so 补位
-
-## 四、upstream 已有、不需要重写的代码
-
-| upstream 文件 | 行数 | 我们的对应文件 | 搬运状态 |
-|--------------|------|---------------|---------|
-| `xllm/core/kernels/ilu/ixformer.h` | 147 | `ex_engine/csrc/ilu/ixformer.h` | ✅ 已搬 |
-| `xllm/core/kernels/ilu/fused_moe.cpp` | 99 | `ex_engine/csrc/ilu_kernel_fused_moe.cpp` | ✅ 已搬 |
-| `xllm/core/layers/ilu/fused_moe.cpp` | 797 | `ex_engine/csrc/ilu_layer_fused_moe.cpp` | ✅ 已搬 |
-| `xllm/core/kernels/ilu/attention.cpp` | 162 | `ex_engine/csrc/ilu_kernel_attention.cpp` | ✅ 已搬 |
-| `xllm/core/layers/ilu/attention.cpp` | 189 | `ex_engine/csrc/ilu_layer_attention.cpp` | ✅ 已搬 |
-| `xllm/core/kernels/ilu/norm.cpp` | 50 | `ex_engine/csrc/ilu_kernel_norm.cpp` | ✅ 已搬 |
-| `xllm/core/kernels/ilu/activation.cpp` | 32 | `ex_engine/csrc/ilu_kernel_activation.cpp` | ✅ 已搬 |
-| `xllm/core/kernels/ilu/rope.cpp` | 31 | `ex_engine/csrc/ilu_kernel_rope.cpp` | ✅ 已搬 |
-| `xllm/core/kernels/ilu/group_gemm.cpp` | 39 | `ex_engine/csrc/ilu_kernel_group_gemm.cpp` | ✅ 已搬 |
-| `xllm/core/kernels/ilu/matmul.cpp` | 73 | `ex_engine/csrc/ilu_kernel_matmul.cpp` | ✅ 已搬 |
-| `xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp` | 576 | `ex_engine/csrc/qwen3_gated_delta_net_base.cpp` | ✅ 已搬 |
-| `ds_vllm/csrc/moe/topk_softmax_kernels.cu` | 874 | `ex_engine/csrc/moe_v055/topk_softmax_kernels.cu` | ✅ 已搬 |
-| `xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh` | ~400 | `ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh` | ✅ 已搬 |
-
-## 五、真机验证 checklist
-
-在真机上按顺序执行:
-
-```bash
-# 1. 验证 ix_moe_bridge.so 编译
-cd /workspace/ex_engine && python3 precompile_ix_bridge.py
-ls build/ix_moe_bridge*.so   # 必须存在
-
-# 2. 验证符号解析
-python3 -c "
-import torch
-import importlib.util
-spec = importlib.util.spec_from_file_location('ix', 'build/ix_moe_bridge.cpython-310-x86_64-linux-gnu.so')
-m = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(m)
-print([x for x in dir(m) if not x.startswith('_')])
-# 应输出: ['topk_softmax', 'moe_gen_idx', 'moe_expand_input', 'moe_group_gemm', 
-#          'silu_and_mul', 'moe_combine_result', 'paged_attention', 'rms_norm',
-#          'fused_add_rms_norm', 'linear', 'reshape_and_cache', 'rotary_embedding']
-"
-
-# 3. 验证 topk_softmax 功能
-python3 -c "
-import torch
-# ... load ix_moe_bridge ...
-gating = torch.randn(4, 64, device='cuda', dtype=torch.float32)
-tw = torch.empty(4, 8, device='cuda', dtype=torch.float32)
-ti = torch.empty(4, 8, device='cuda', dtype=torch.int32)
-tei = torch.empty(4, 8, device='cuda', dtype=torch.int32)
-m.topk_softmax(tw, ti, tei, gating)
-print('topk_weights:', tw)
-print('topk_ids:', ti)
-"
-
-# 4. 验证 GDN 不再 NaN
-# (需要先修复 GDN prefill 代码)
-
-# 5. 启动服务验证
-python3 -m vllm.entrypoints.openai.api_server --model /model ...
-```
-
-## 六、最关键发现：07-23 的 base image 自带完整 corex_* chain
-
-**07-23 日志证据** (dockerrizhi.txt):
-```
-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_moe.py:339  → Using CoreX fused MoE prefill operator: tokens=4096, kernel=expert-grouped-wmma ✅
-corex_fa2.py:333  → Using CoreX FA2 packed prefill: B=2 Hq=4 Hkv=1 D=256 ✅
-corex_fa2.py:507  → Using CoreX paged FA2 chunked prefill ✅
-```
-
-**08-07 日志**: 零条 corex_* 加载记录。取而代之的是 `qwen3_5.py:445 NaN in prefill` + `_custom_ops.py:58 topk_softmax not found`。
-
-**根因**: 08-07 提交部署了我们自己的 `qwen3_5.py`，覆盖了 base image 自带的版本，打断了 `corex_gdn.py` / `corex_moe.py` / `corex_fa2.py` 的调用链。
-
-**当前状态**: `patch_ops.sh v2` 已经有条件跳过逻辑（`_QW_SIZE > 1000 → KEEPING IT`），但需要确保下次提交时不再触发 qwen3_5.py 覆盖。
-
-**结论**: 如果 base image 有工作的 corex_* chain，我们只需要:
-1. 不覆盖 qwen3_5.py
-2. 只部署 serving 层（protocol/serving_chat/api_server/tool_parser）
-3. `max_model_len=80000`（已修）
-4. `ix_moe_bridge.so` 作为备用（如果 base 的 _custom_ops 有路径碰到 topk_softmax）
-
-## 七、代码量评估
-
-| 组件 | 文件数 | 总行数 | 状态 |
-|------|--------|--------|------|
-| ex_engine/csrc (C++) | 39 | ~8000 | 全部已有，需真机编译 |
-| ex_engine/python (Python) | 7 | ~1200 | 全部已有，dispatch chain 完整 |
-| qwen3_6_scripts (serving) | 20+ | ~6000 | 全部已有，patch_ops.sh 管部署 |
-| upstream_ref (xllm reference) | 500+ | ~100K | 参考用，关键文件已搬到 ex_engine |
-
-**结论**: 代码量是够的。问题不是代码不够，而是:
-1. GDN NaN 没修（需要用 xllm 的 fp32 accumulation 逻辑替换）
-2. ix_moe_bridge.so 从未在真机编译成功
-3. 没有 "不允许 fallback" 的硬要求落实到代码里

COMPETITIVE_ANALYSIS_AND_FIX_PLAN.md

@@ -1,121 +0,0 @@
-# 竞赛对比分析 & 修复计划
-
-## 一、核心数据对比
-
-| 模块 | 对手 Sub168 | 我们 Sub508 | 差距 |
-|------|-----------|-----------|------|
-| **functional** | 48/52 PASS (92.3%) | 21/51 PASS (41.2%) | **-51%** |
-| **case_truncation** | score=1.0 (8192 tokens输出完整) | score=0.0 (引擎崩溃) | **致命** |
-| **replay_tencent** | score=60194 (94/881成功,tps avg 11.86) | score=0.0 (881/881 connection refused) | **致命** |
-| **opencompass** | 0.0 (server也崩了) | 0.0 (同上) | 平 |
-| **总分** | **60194.6** | **0.0** | -- |
-
-## 二、Sub508 崩溃根因链
-
-```
-t2_n_2 (n=2请求) → get_scheduler_config() 异常 → 引擎进程死亡
-→ 后续所有请求 Connection Refused → 30个FAIL级联
-→ case_truncation/replay/opencompass 全部0分
-```
-
-**关键事实：t2_n_2 崩溃发生在 06:42:45，之后所有模块都是在引擎已死的情况下跑的。**
-
-## 三、对手 Sub168 的弱点（我们已经修复的）
-
-1. **`max_completion_tokens` 被拒** — 对手 `extra="forbid"` 导致 replay 中所有带此字段的请求返回 400。我们已添加该字段到 protocol.py，replay 中不会被拒。
-2. **`tool_calls` content=None 被拒** — 对手的 replay preflight 失败（"Each message must have at least one of 'content' or 'reasoning_content'"）。我们已修复 chat_utils.py 中 content=None 的处理。
-3. **d06_cache_hit FAIL** — 对手没有 prefix caching，我们 PASS。
-4. **t3_max_tokens_1/64/max 3个FAIL** — 对手也有3个max_tokens测试失败。
-
-**对手 replay 中 787/881 失败(89.3%)，只有 94 个成功。我们的目标是超越这个。**
-
-## 四、我们需要修复的问题（按优先级排序）
-
-### P0 — 引擎稳定性（决定能否拿分的前提）
-
-| 问题 | 根因 | 修复位置 |
-|------|------|----------|
-| **t2_n_2 → 引擎崩溃级联** | `get_scheduler_config()` 异常 + n>1 未处理 | `qwen3_6_scripts/serving_chat.py` + `protocol.py` |
-| **引擎OOM死亡** | 单个长请求耗尽GPU内存后整个进程死 | 需要在 worker/model_runner.py 加 OOM catch |
-
-已有 commit 修复（994c657 clamp n>1, c241764 try-catch scheduler），但 **Sub508 用的是修复前的代码**。Sub509 日志确认 d01 能跑（95.85s），但 d03 仍然 FAIL。
-
-### P1 — d03_tool_call FAIL（功能测试核心分）
-
-**Sub508**: `tools=0 finish=stop reasoning[0]` (49.04s)
-**Sub509**: `tools=0 finish=stop reasoning[0]` (49.04s)
-**对手**: `tool=get_weather args="{'city': 'Beijing'}" finish=tool_calls` (2.12s)
-
-**根因分析**：
-- 对手 d03 只用了 2.12s，模型直接输出 tool_call XML，tool parser 正确解析
-- 我们用了 49.04s，模型在 thinking 中耗尽了时间，没有产生 `<tool_call>` 标签
-- commit e0344b1 说"禁用 tool_call 请求的 thinking"，但 Sub509 的 d03 仍显示 `reasoning[0]`
-- **真正的问题**：当 `tool_choice=auto` 且有 tools 时，需要在 chat_template 中设置 `enable_thinking=False`，否则 Qwen3 会先 think 再输出，大量token浪费在思考上
-
-**修复方案**：在 `serving_chat.py` 的 `create_chat_completion` 中，当检测到 `request.tools` 且 `tool_choice != "none"` 时，在 `chat_template_kwargs` 中注入 `enable_thinking=False`。
-
-### P1 — d05_multimodal HTTP 400
-
-对手 PASS (content[374])，我们 HTTP 400。
-可能是多模态请求格式/图片解码问题。需要检查 chat_utils.py 的图片处理路径。
-
-### P1 — d07_reasoning_plus_content
-
-对手 PASS (reasoning[3489] content[962])，我们 FAIL (reasoning[131] content[0])。
-模型 think 后不产生 content。这是模型行为问题，但可以通过调低 thinking budget 或调整 temperature 来缓解。
-
-### P2 — t1a_thinking_true / t1c_thinking_default
-
-对手 PASS (reasoning[541] / [411])，我们 FAIL (reasoning[0])。
-**根因**：模型在短回答场景下不触发 thinking。可能需要在 chat_template 中确保 `enable_thinking=True` 是默认值。检查 Qwen3.6 的 chat_template 是否正确注入了 `<think>` 标签。
-
-### P2 — d10_thinking_disable_ctk 乱码输出
-
-对手输出 `'4'`（正确），我们输出乱码 `"presت< **sama一..."`。
-模型在 thinking disabled 模式下输出质量极差。这是模型+chat_template 的交互问题。
-
-### P3 — 速度差距
-
-| 测试 | 对手 | 我们 | 倍数 |
-|------|------|------|------|
-| d01 | 8.49s | 95.85s | **11x慢** |
-| d04 | 17.78s | 128.74s | **7x慢** |
-| d03 | 2.12s | 49.04s | **23x慢** |
-
-速度问题核心：BI-V100 硬件本身比 NVIDIA GPU 慢，但 10x 的差距说明还有架构问题。对手的 output_tps 平均 11.86，decode 阶段 tps 在 2.4-22.7 之间。
-
-## 五、修复代码的具体文件
-
-需要修改的文件（全部在 `qwen3_6_scripts/` 中，会被 patch_ops.sh 部署）：
-
-1. **`serving_chat.py`** — tool_call 时注入 `enable_thinking=False`
-2. **`protocol.py`** — 确认 `extra="forbid"` 已经去掉（已做），确认 `thinking` 字段被正确传递
-3. **`chat_utils.py`** — 多模态请求处理、content=None 容错
-4. **`model_runner.py`** — OOM recovery
-5. **`qwen3_5.py`** — 检查模型是否正确处理 `enable_thinking` 参数
-6. **`computility-run.yaml`** — 考虑调整 `--max-num-seqs` / `--gpu-memory-utilization`
-
-## 六、对手的 replay 得分结构
-
-对手 881 个请求中：
-- 94 个成功 (10.7%)
-- 77 个因 `max_completion_tokens` extra_forbidden 而 400
-- 704 个 connection refused（server也崩了！）
-- output_tps_avg = 11.86, output_tps_p50 = 12.97
-
-**关键发现：对手的 server 也在 replay 后期崩溃了（704 个 connection refused）。但他在崩溃前完成了 94 个请求。**
-
-我们的优势：
-- 我们已修复 `max_completion_tokens` → 对手的 77 个 400 我们不会有
-- 我们已修复 `tool_calls content=None` → 对手的 tool preflight fail 我们不会有
-- 我们有 prefix caching → 对手没有
-
-**如果我们能保持引擎稳定不崩溃，仅靠不拒绝 max_completion_tokens 的请求，就能多处理 77+ 个请求，超过对手。**
-
-## 七、下一步行动
-
-1. 修复 `serving_chat.py`：tool_call 时禁用 thinking
-2. 确认 n>1 clamp 和 scheduler try-catch 在 patch 文件中生效
-3. 测试 OOM 恢复逻辑
-4. 调整 computility-run.yaml 参数确保稳定性
-5. 提交部署，跑测试

DLOPEN_DEV_PLAN.md

@@ -0,0 +1,155 @@
+# dlopen SO开发计划 — 从日志到代码
+
+> 基于 comp168 docker (2d5232c5) 日志分析 + 真机代码 tree (不带 --depth)
+> 原则：upstream已有的搬过来，接口对上，不允许fallback，不允许全新开发
+
+---
+
+## 一、真机调用链现状（qwen3_5.py imports）
+
+qwen3_5.py 声明了 **11个** corex SO模块的 import：
+
+| # | 模块名 | prebuilt .so | .cu源码 | build脚本 | qwen3_5.py调用点 | 状态 |
+|---|--------|-------------|---------|-----------|-----------------|------|
+| 1 | corex_gdn_causal_conv | ✅ | ✅ | ✅ | L1158: conv更新 | **就绪** |
+| 2 | corex_gdn_gated_norm | ✅ | ✅ | ✅ | L848: 反向norm | **就绪** |
+| 3 | corex_gdn_beta_decay | ✅ | ✅ | ✅ | L1215: 衰减计算 | **就绪** |
+| 4 | corex_gdn_qk_map | ✅ | ✅ | ✅ | L1258: QK映射 | **就绪** |
+| 5 | corex_gdn_packed_decode | ✅ | ✅ | ✅ | L1195: 打包解码 | **就绪** |
+| 6 | corex_attn_head_rms_norm | ✅ | ✅ | ✅ | L1322: 头归一化 | **就绪** |
+| 7 | corex_moe_exact_reduce | ✅ | ✅ | ✅ | L1707: MoE精确归约 | **就绪** |
+| 8 | corex_moe_weight_gather | ✅ | ✅ | ✅ | L1681: 权重收集 | **就绪** |
+| 9 | corex_moe_direct_routed | ✅ | ✅ | ✅ | L1659: 直接路由MoE | **就绪** |
+| 10 | corex_moe_topk_softmax | ✅ | ✅ | ✅ | L1621: topk+softmax | **就绪** |
+| 11 | corex_moe_index_combine | ❌ 无prebuilt | ✅ | ✅ | L1719: 索引合并 | **需在docker build编译** |
+
+## 二、prebuilt有但qwen3_5.py没引用的SO
+
+| 模块名 | prebuilt | .cu源码 | qwen3_5.py引用 | 说明 |
+|--------|---------|---------|---------------|------|
+| corex_block_major_kv_transfer | ✅ | ✅ | ❌ | block_major_kv_cache.py用 |
+| corex_fused_paged_prefill | ✅ | ✅ (split4版) | ❌ | paged_attn.py用 |
+| corex_paged_kv_gather | ✅ | ✅ | ❌ | paged_attn.py用 |
+
+## 三、有.cu但无prebuilt的模块
+
+| 模块名 | .cu源码 | 说明 | 行动 |
+|--------|---------|------|------|
+| corex_gdn_chunk_recurrent | ✅ (10807字节) | GDN prefill chunked recurrent | **需precompile，可能是NaN修复的关键** |
+| corex_fused_paged_prefill_split4 | ✅ (20172字节) | 分4路prefill attention | prebuilt有 corex_fused_paged_prefill (名字不同) |
+| corex_moe_index_combine | ✅ (5554字节) | patch_ops.sh已有编译步骤 | **Docker内编译** |
+| corex_query_tiled_paged_prefill | ✅ (20409字节) | Q-tiled prefill | 当前paged_attn.py的Python版替代 |
+
+## 四、comp168日志揭示的关键差距
+
+comp168（竞争对手sub168）的Docker工作正常：
+- GDN：用 corex_gdn.so 的fused kernel，**无NaN**
+- MoE：用自己的 topk_softmax 实现 + WMMA group_gemm，**不依赖 ixf_F.vllm_moe_topk_softmax**
+- 权重：17.35 GB（我们16.23 GB）
+- model_runner.py: 用base镜像原版(1074行)，不是我们的1119行版
+
+我们的Docker(sub655)的问题：
+- GDN：99.98% NaN → nan_to_num → 输出垃圾
+- MoE：fallback到PyTorch loop → 约50x慢
+- 服务器最终崩溃 → Connection refused → 881个replay请求全失败
+
+## 五、现在的代码量够不够？
+
+```
+qwen3_6_scripts/
+├── 15个 corex_*.cu 文件       (总计 ~115K 字节 CUDA源码)
+├── 14个 build_corex_*.sh      (编译脚本)
+├── 13个 prebuilt/*.so          (已编译二进制)
+├── qwen3_5.py                  (1700+行，模型实现)
+├── patch_ops.sh                (部署脚本)
+├── paged_attn.py               (paged attention)
+├── serving_chat.py + protocol.py + api_server.py (serving层)
+├── vendor_overrides/           (vllm核心override，6文件)
+└── ...
+
+ex_engine/
+├── csrc/                       (C++ bridge代码，24个文件)
+├── python/                     (Python bridge代码，7个文件)
+├── xllm_kernels/               (xllm上游kernel，8个文件)
+└── xllm_layers/ + xllm_models/ (xllm上游层/模型实现)
+
+upstream_ref/
+├── ds_vllm/                    (最新vllm参考实现)
+├── xllm/                       (xllm完整参考)
+├── fla/                        (flash-linear-attention参考)
+└── vllm_gdn/                   (vllm GDN参考实现)
+```
+
+**回答你的问题：代码数量是够的。** 15个.cu、13个prebuilt .so、qwen3_5.py已经完整引用了所有11个import。问题不是代码数量，是：
+
+1. **corex_moe_index_combine.so 没有prebuilt** — 需要在docker build时在线编译
+2. **corex_gdn_chunk_recurrent.so 没有prebuilt** — 10K字节的GDN prefill kernel，可能是解决NaN的关键
+3. **patch_ops.sh 只编译了 moe_index_combine** — 其余12个走prebuilt安装
+
+## 六、下一步行动（代码开发，不是推理）
+
+### 立即要做的3件事：
+
+**1. 把 corex_gdn_chunk_recurrent 加入 prebuilt 或 patch_ops.sh 编译链**
+
+这个.cu存在（10807字节），build脚本也存在，但既没有prebuilt .so，也没在patch_ops.sh里编译。真机上需要：
+
+```bash
+# 在你的BI-V100真机上：
+cd /home/dylan/project_6/qwen3_6_scripts
+bash build_corex_gdn_chunk_recurrent.sh /usr/local/corex/lib/python3/dist-packages/vllm
+# 如果成功，把.so拷到 prebuilt/corex-3.2.3-ivcore10/
+```
+
+**2. qwen3_5.py GDN prefill路径需要对接 chunk_recurrent kernel**
+
+当前qwen3_5.py的GDN prefill fallback是纯PyTorch `_torch_chunk_gated_delta_rule`，产生NaN。corex_gdn_chunk_recurrent.cu 是 fp32 accumulation 的 kernel — 应该能解决NaN。需要在qwen3_5.py里加上对应的 import + dispatch。
+
+**3. 把 corex_fused_paged_prefill_split4.cu precompile**
+
+这个20K字节的kernel对应prefill attention加速，prebuilt目录有 `corex_fused_paged_prefill.so`（可能是同一个的改名），需要确认对应关系。
+
+### 在真机上验证步骤：
+
+```bash
+# 单卡验证：
+cd /home/dylan/project_6
+python3 -c "
+import torch
+# 测试prebuilt SO能否加载
+import importlib.util
+spec = importlib.util.spec_from_file_location('corex_gdn_causal_conv', 
+    'qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/corex_gdn_causal_conv.so')
+mod = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(mod)
+print('corex_gdn_causal_conv loaded:', dir(mod))
+"
+```
+
+## 七、commit 9ff2450（能得分的版本）
+
+这个commit不在当前仓库里。你说它是 `clean: remove build artifacts from docker context`，date Aug 12 07:58。这意味着它是在current HEAD (17fdf7e2) 之后的commit，可能在另一个branch或还没push。
+
+**需要你执行：**
+```bash
+git log --all --oneline | grep 9ff2450
+# 或者
+git push origin main  # 如果在真机上有unpushed commits
+```
+
+## 八、ex_engine upstream搬运清单
+
+ex_engine里有大量代码但 **没有接入 patch_ops.sh 部署链**。以下是已有但未使用的：
+
+| 文件 | 功能 | upstream来源 | 接入状态 |
+|------|------|-------------|---------|
+| ex_engine/python/corex_gdn.py | GDN完整dispatch | 自己写的 | ❌ 未部署 |
+| ex_engine/python/corex_moe.py | MoE完整dispatch | 自己写的 | ❌ 未部署 |
+| ex_engine/python/ix_bridge.py | C++→Python bridge | 自己写的 | ❌ 未部署 |
+| ex_engine/csrc/ix_full_bridge.cpp | ixformer C++桥 | 基于symbol probe | ❌ 未部署 |
+| ex_engine/xllm_kernels/cuda/moe/*.cu | MoE CUDA kernels | xllm upstream | ❌ 未部署 |
+| ex_engine/xllm_layers/npu_torch/*.cpp | 层实现 | xllm upstream | ❌ 未部署 |
+
+**这些不需要重写，但接口要对上后再搬。** 特别是 ix_full_bridge.cpp 里明确说了 "MoE functions are NOT in base image"，所以 MoE 必须走 prebuilt .so + Python fallback 路线，而不是试图 dlopen 不存在的 ixformer MoE symbols。
+
+现在的策略（13个prebuilt .so + 1个在线编译）已经是正确的路线。

ENGINEX_INJECTION_MAP.md

@@ -1,94 +0,0 @@
-# EngineX vllm Injection Point Map
-
-> **Source**: `enginex-vllm-bi100-qwen36-main.zip` (101MB, 1444 files)
-> **Generated**: 2026-08-02 from full source analysis
-
----
-
-## 关键发现
-
-### 1. 不是 C++ CUDA 文件注入 — 是 Python 层
-
-EngineX vllm 的 CUDA kernels 全部预编译在 `ixformer.functions` (ixf_F) 中，打包在基础镜像里。
-`_custom_ops.py` 是 Python 薄封装层，调用 `ixf_F.vllm_single_query_cached_kv_attention()` 等。
-
-**没有 .cu 文件可以直接 patch。** muh 的 gen_patch.py 需要改为 patch Python 文件，不是 C++ 文件。
-
-### 2. paged_attention_v2 未实现
-
-```python
-def paged_attention_v2(...) -> None:
-    raise NotImplementedError()
-```
-
-且 `use_v1 = True` 硬编码覆盖了启发式逻辑。所有 decode 都走 v1。
-
-### 3. 实际可调参数 (THE TUNING SURFACE)
-
-| 参数 | 文件 | 当前值 | 作用 | 优先级 |
-|------|------|--------|------|--------|
-| `_PARTITION_SIZE` | `vllm/attention/ops/paged_attn.py:13` | 512 | PagedAttention partition (v2 用) | 低 (v2 disabled) |
-| `use_v1` | `paged_attn.py:128` | `True` (hardcoded) | 强制 v1 | **P0** — 解锁 v2 可能提升长序列 |
-| `BLOCK` | `prefix_prefill.py:712` | 128 (cc≥80) / 64 | Triton prefill tile size | **P0** — 直接影响 Input TPS |
-| `NUM_WARPS` | `prefix_prefill.py:713` | 8 | Triton warp count | **P0** |
-| `BLOCK_SIZE_M/N/K` | `fused_moe.py:342-344` | 64/64/32 | MoE kernel tile | **P0** — Qwen3.6 是 MoE |
-| `get_max_shared_memory` | `_custom_ops.py:892` | `32 * 1024` | SMEM 上限声明 | **P0** — 可能错误限制性能 |
-| Triton flash attention configs | `triton_flash_attention.py:214-303` | 8 个 triton.Config | Triton autotune 搜索空间 | P1 |
-
-### 4. SMEM 32KB vs 48KB 冲突
-
-`_custom_ops.py:892` 返回 `32 * 1024` (32KB)。
-但 `hardware.cuh` 和 muh 假设 49152 (48KB)。
-如果 BI-V100 实际 SMEM 是 32KB，则 muh 所有 tuning 的 SMEM 约束都需要从 48KB 降到 32KB。
-
-### 5. ixf_F kernel 列表 (不可改，只能调参)
-
-| Python 封装 | ixf_F 调用 | 说明 |
-|-------------|-----------|------|
-| `paged_attention_v1` | `ixf_F.vllm_single_query_cached_kv_attention` | decode 核心 |
-| `silu_and_mul` | `ixf_F.silu_and_mul` | SwiGLU 激活 |
-| `rms_norm` | `ixf_F.rms_norm` | LayerNorm |
-| `fused_add_rms_norm` | `ixf_F.fused_add_rms_norm` | 融合残差+norm |
-| `rotary_embedding` | `ixf_F.vllm_rotary_embedding_neox` | RoPE 位置编码 |
-| `reshape_and_cache` | `ixf_F.vllm_cache_ops_reshape_and_cache` | KV cache 写入 |
-| `copy_blocks` | `ixf_F.copy_blocks` | prefix cache block 复制 |
-| `moe_align_block_size` | `ixf_F.vllm_moe_align_block_size` | MoE token 排列 |
-| `invoke_fused_moe_kernel` | `ixf_F.vllm_invoke_fused_moe_kernel` | MoE GEMM |
-| `topk_softmax` | `ixf_F.vllm_moe_topk_softmax` | MoE routing |
-| `cutlass_scaled_mm` | `ixf_F.w8a8` | INT8 矩阵乘 |
-
-### 6. Triton kernels (可直接修改)
-
-这些是 Python Triton JIT 编译的 kernel，可以直接改源码：
-
-- `prefix_prefill.py` — 3 个 `_fwd_kernel` 变体 (context attention)
-- `triton_flash_attention.py` — Triton flash attention (8 个 autotune configs)
-- `fused_moe.py` — MoE GEMM kernel (Triton, 自定义 config)
-
----
-
-## muh 策略修正
-
-### 旧策略 (假设 C++ injection)
-```
-CCCL tuning_*.cuh → muh bi100_* → gen_patch.py → C++ #define 注入 → 编译 .so
-```
-
-### 新策略 (实际 Python injection)
-```
-层1: Python 参数调优
-  paged_attn.py: _PARTITION_SIZE, use_v1
-  prefix_prefill.py: BLOCK, NUM_WARPS
-  fused_moe.py: BLOCK_SIZE_M/N/K
-  _custom_ops.py: get_max_shared_memory (32KB→实测值)
-
-层2: Triton kernel 优化
-  prefix_prefill.py: 3 个 _fwd_kernel — tile size, loop structure
-  triton_flash_attention.py: autotune config 添加 BI-V100 特化
-  fused_moe.py: MoE GEMM kernel tune
-
-层3: CCCL/muh 知识迁移
-  用 CCCL 的 tuning 方法论指导 Triton kernel 参数选择
-  不是直接注入 C++ 值，而是把 CCCL 的 policy_selector 逻辑
-  翻译成 Triton constexpr 参数
-```

ENGINE_CODEPATH_TIMELINE.md

@@ -1,150 +0,0 @@
-# Engine Code Path Timeline: Sub168 vs Our Sub508/509
-
-**Purpose**: Anyone reading this repo can understand the exact runtime difference in 2 minutes instead of re-deriving from raw logs.
-
-## 1. Boot Sequence Comparison
-
-```
-TIME          SUB168 (07-23, score=60194)              OUR SUB508 (08-07, score=0)
-──────────────────────────────────────────────────────────────────────────────────
-+0s           api_server.py:530 → vLLM 0.6.3           api_server.py:530 → vLLM 0.6.3
-              max_model_len=256000                      max_model_len=256000 (same)
-              max_num_seqs=2, gpu_mem=0.95              max_num_seqs=2, gpu_mem=0.95 (same)
-              chunked_prefill=True                      chunked_prefill=True (same)
-
-+10s          model_runner.py:1074 load start           model_runner.py:1119 load start
-              ↑ DIFFERENT line number                   ↑ DIFFERENT line number
-              ↑ (base image native model_runner)        ↑ (our patched model_runner)
-
-+18s          weights = 17.3529 GB                      weights = 16.2303 GB
-              ↑ 1.1GB MORE (corex state buffers)        ↑ 1.1GB LESS (no corex buffers)
-
-+180s         corex_gdn.py:56  → load libcorex_gdn.so  qwen3_5.py:445 → NaN in prefill layer 0
-              corex_gdn.py:228 → GDN prefill OK         ↑ PyTorch GDN produces NaN (99.98%)
-              corex_moe.py:339 → MoE prefill OK         qwen3_5.py:913 → FusedMoE FAILED
-              corex_fa2.py:333 → FA2 prefill OK         ↑ ixformer.functions missing topk_softmax
-              ↑ ALL THREE CoreX accelerators loaded     ↑ ZERO accelerators, all fallback
-
-+182s         GPU blocks: 19259                         GPU blocks: ~19000 (similar)
-              Ready to serve                            Ready to serve (but 10x slower)
-```
-
-## 2. Call Chain During Inference
-
-### Sub168 (with CoreX) — d01_basic_nostream: 8.49s
-```
-serving_chat.py → create_chat_completion()
-  → engine.generate()
-    → model_runner.py:1074 execute_model()
-      → qwen3_5.py:1421 Qwen3_5ForCausalLM.forward()
-        → qwen3_5.py:1165 Qwen3_5Model.forward()  (decoder layers loop)
-          → qwen3_5.py:1086 Qwen3_5DecoderLayer.forward()
-            ├─ GatedDeltaNet layers (4 of 36):
-            │   ├─ PREFILL: corex_gdn.py:228 → libcorex_gdn.so (fused CUDA kernel)
-            │   └─ DECODE:  corex_gdn.py:138 → libcorex_gdn.so (fused CUDA kernel)
-            ├─ MoE layers (all 36):
-            │   ├─ PREFILL: corex_moe.py:339 → libcorex_moe.so (expert-grouped-wmma)
-            │   └─ DECODE:  corex_moe.py:249 → libcorex_moe.so (fused MoE decode)
-            └─ Attention (32 of 36 layers):
-                ├─ PREFILL: corex_fa2.py:333 → libcorex_fa2.so (packed FA2)
-                └─ DECODE:  corex_fa2.py:225 → libcorex_fa2.so (paged decode)
-```
-
-### Our Sub508 (no CoreX) — d01_basic_nostream: 95.87s (11.3x slower)
-```
-serving_chat.py → create_chat_completion()
-  → engine.generate()
-    → model_runner.py:1119 execute_model()
-      → qwen3_5.py:1369 Qwen3_5ForCausalLM.forward()  (52 lines shorter!)
-        → qwen3_5.py:???? Qwen3_5Model.forward()
-          → qwen3_5.py:???? Qwen3_5DecoderLayer.forward()
-            ├─ GatedDeltaNet layers (4 of 36):
-            │   ├─ PREFILL: pure PyTorch conv1d → matmul → softmax (NaN!)
-            │   └─ DECODE:  pure PyTorch _torch_causal_conv1d_update
-            ├─ MoE layers (all 36):
-            │   ├─ PREFILL: PyTorch loop over unique_eids (SLOW)
-            │   └─ DECODE:  PyTorch batched GEMM fallback
-            └─ Attention (32 of 36 layers):
-                ├─ PREFILL: xformers _run_sdpa_fallback (patched, matmul+softmax)
-                └─ DECODE:  xformers _run_sdpa_fallback
-```
-
-## 3. The Crash Chain (Sub508/509 → Score 0)
-
-```
-FUNCTIONAL TEST SEQUENCE:
-d01_basic_nostream  ✓ PASS (95.87s — slow but works)
-d02_stream_usage    ✓ PASS (1.84s)
-d03_tool_call       ✗ FAIL (49.04s — model thinks instead of emitting tool XML)
-d04_reasoning       ✓ PASS (128.74s)
-   ... more tests pass ...
-t2_n_2              ✗ FAIL → HTTP 500 → ENGINE PROCESS DIES
-   ↓
-t3_max_tokens_none  ✗ FAIL → HTTP 500 (engine dead, Connection Refused)
-t3_max_tokens_1     ✗ FAIL → HTTP 500
-t3_max_tokens_64    ✗ FAIL → HTTP 500
-   ... 25 more tests ...
-t16c_empty_messages ✗ FAIL → HTTP 500
-───────────────────────────────────────
-functional score: 21/51 = 0.412 (passed before crash)
-
-case_truncation     → Connection Refused → score=0.0
-replay_tencent      → 881/881 Connection Refused → score=0.0
-opencompass          → Connection Refused → score=0.0
-───────────────────────────────────────
-TOTAL: 0.0  (engine was dead for 90% of evaluation)
-```
-
-## 4. CoreX Dispatch Gap — The 52-Line Difference
-
-Sub168's qwen3_5.py has ~1421 lines. Ours has 1369.
-The missing ~52 lines are CoreX dispatch wrappers:
-
-```python
-# WHAT SUB168 HAS (reconstructed from log evidence):
-
-# In GatedDeltaNet.__init__:
-try:
-    from vllm.model_executor.models.corex_gdn import CoreXGDN
-    self._corex_gdn = CoreXGDN(...)  # loads libcorex_gdn.so
-except ImportError:
-    self._corex_gdn = None
-
-# In GatedDeltaNet.forward() prefill path:
-if self._corex_gdn is not None:
-    result = self._corex_gdn.prefill(...)  # → corex_gdn.py:228
-else:
-    result = self._pytorch_prefill(...)    # our current pure PyTorch
-
-# In Qwen3_5MoE.forward():
-try:
-    from vllm.model_executor.models.corex_moe import corex_moe_forward
-    result = corex_moe_forward(...)        # → corex_moe.py:339
-except:
-    result = self._pytorch_moe_forward(...)  # our current loop
-```
-
-## 5. Environment Variables (already set in YAML)
-
-```yaml
-VLLM_COREX_GDN_LIBRARY: /usr/local/corex/lib64/libcorex_gdn.so
-VLLM_COREX_MOE_LIBRARY: /usr/local/corex/lib64/libcorex_moe.so
-VLLM_COREX_FA2_LIBRARY: /usr/local/corex/lib64/libcorex_fa2.so
-```
-
-These .so files exist in the base image. The Python wrappers
-(`corex_gdn.py`, `corex_moe.py`, `corex_fa2.py`) also exist in
-the base image at:
-`/usr/local/corex/lib/python3/dist-packages/vllm/model_executor/models/`
-
-**Our qwen3_5.py simply never imports them.**
-
-## 6. What Needs To Happen
-
-Add try/except CoreX dispatch in 3 places in qwen3_5.py:
-1. `GatedDeltaNet.forward()` — prefill + decode paths
-2. `Qwen3_5MoE.forward()` — prefill + decode MoE dispatch  
-3. Attention — already handled by xformers patches (corex_fa2 is separate)
-
-CCCL pattern: `dispatch_with_env` — try native kernel first, fallback on error.
-Our Python equivalent: `try: corex_forward() except: pytorch_forward()`

GROUND_TRUTH_STATUS.md

@@ -1,124 +0,0 @@
-# project_6 真实状态报告
-
-生成时间: 2026-08-05, commit 96f6465
-
-## 一句话总结
-
-**enginex 没有 .cu 源码，gen_patch 的 C++ injection 管道全部失效。** 实际可用的优化路径只有 Python/Triton 层面的参数调优。muh 的 27 个 C++ tuning headers 是正确的架构设计，但在竞赛引擎上无处注入。
-
----
-
-## 1. 竞赛引擎的致命事实
-
-```
-gen_patch.py 第 47 行:
-  WARNING: ALL csrc/*.cu targets are DEAD — files do not exist.
-  enginex-vllm-bi100-qwen36 ships: Python + precompiled .so + Triton.
-  No .cu source files. gen_patch patches have zero effect.
-```
-
-enginex 交付物 = Python 文件 + 预编译 .so + Triton kernels。  
-不提供 C 源码 → 无法修改 CUDA kernel → C++ tuning header 无法注入到 vllm 的编译产物里。
-
-**真正的优化路径:**
-- Triton kernels (prefix_prefill.py, paged_attn.py): 可以改 BLOCK、NUM_WARPS 等 JIT 参数
-- Python 配置层 (computility-run.yaml): max_model_len、gpu_memory_utilization 等
-- 模型适配 (qwen3_5.py): MoE routing、attention 实现
-
-## 2. 已有的 benchmark 数据 (真实的)
-
-| 算法域 | 已跑配置数 | 来源 |
-|--------|-----------|------|
-| flash_attn | 22 configs | bi100_configs.json, SMEM 约束扫描 |
-| prefill (Triton) | 9 configs | bi100_configs.json, BLOCK×NUM_WARPS |
-| MoE | 5 configs | bi100_configs.json, BLOCK_SIZE_M |
-| reduce/scan/topk CUB | 0 | bench_bi100.py 已写但需要 BI-V100 硬件才能跑 |
-
-## 3. muh C++ headers vs CCCL 覆盖率
-
-| 算法 | muh 行数 | CCCL 行数 | 覆盖率 | 竞赛优先级 |
-|------|---------|---------|--------|-----------|
-| reduce | 297 | 478 | 62% | **P0** — Output TPS 83% 权重 |
-| scan | 352 | 1525 | 23% | **P0** — softmax 累积 |
-| topk | 113 | 121 | 93% | **P0** — sampling 路径 |
-| transform | 185 | 549 | 33% | P1 — RMSNorm/SiLU |
-| select_if | 459 | 2729 | 16% | P1 — token filtering |
-| radix_sort | 222 | 2381 | 9% | P1 — full sort path |
-| scan_by_key | 145 | 2008 | 7% | P1 — per-seq softmax |
-| reduce_by_key | 171 | 1735 | 9% | P1 — score aggregation |
-| unique_by_key | 166 | 1539 | 10% | P1 — KV cache dedup |
-| 其余 18 个 | 33-189 | 78-788 | 10-65% | P2 |
-
-总计: muh 3618 行 vs CCCL 17000+ 行 = 平均 21% 覆盖率
-
-## 4. CCCL 资产完整性
-
-cccl_upstream/ 34MB, 3432 files — 是精选提取, 不是 full clone。
-
-**已有 (竞赛必需的全有):**
-- 27/27 tuning headers ✓
-- 32/32 dispatch implementations ✓ 
-- 25/25 agent kernels ✓
-- 60/60 Thrust examples ✓
-- 243 CUB tests ✓
-- 78 CUB benchmark .cu files ✓
-- 230 Thrust tests ✓
-- 48 Thrust benchmark algorithms ✓
-
-**不需要 full clone。** 缺的 ~21000 文件是 CI/CD、cudax、Python bindings、docs。
-
-## 5. 真正的行动路径
-
-### 短期 (功能测试通过)
-竞赛门控: 50+ 功能测试全通过 + 效果偏差 ≤ ±4%
-
-关键文件:
-- `computility-run.yaml` — 控制 vllm 启动参数
-- `qwen3_6_scripts/qwen3_5.py` (588行) — MoE 模型适配
-- `prefix_prefill.py` — Triton prefill kernel, 可调 BLOCK/NUM_WARPS
-- `paged_attn.py` — Triton decode kernel
-
-### 中期 (性能优化)
-目标: Token 吞吐加权值 ≥ 8000
-
-```
-加权值 = Output_TPS × 16.796 + Input_TPS × 2.799 + Cache_TPS × 0.56
-```
-
-**Output TPS (83%):** decode kernel → paged_attn.py Triton 参数优化  
-**Input TPS (14%):** prefill kernel → prefix_prefill.py Triton 参数优化  
-**Cache TPS (3%):** prefix caching 配置
-
-### 长期 (如果能编译 C++)
-如果能获取 EngineX 的 C 编译环境:
-- muh C++ headers 可以直接注入
-- bench_bi100.py 的 CUB parameter sweep 可以在 BI-V100 上跑
-- 这条路 ROI 最高但依赖竞赛方提供编译链
-
-## 6. 代码架构
-
-```
-project_6/
-├── computility-run.yaml          ← 竞赛提交配置 (直接影响评测)
-├── baseline.muh                  ← muh 格式的 vllm 配置
-├── Dockerfile                    ← 竞赛镜像构建
-├── cccl_upstream/                ← CCCL 精选 (34MB, 3432 files)
-│   ├── cub/                      ← CUB: dispatch/tuning/agent/test/bench
-│   ├── thrust/                   ← Thrust: examples/testing/benchmarks
-│   └── libcudacxx/               ← CUDA 标准库
-├── muh/                          ← kernel tuning 框架 (544KB)
-│   ├── include/muh/tuning/       ← 27 个 BI-V100 tuning headers
-│   ├── bench_bi100.py            ← CUB parameter sweep runner
-│   ├── gen_patch.py              ← vllm patch 生成 (C++ 注入点已死)
-│   ├── gen_yaml.py               ← computility-run.yaml 生成
-│   └── parse.py                  ← .muh 配置解析器
-├── muh_kernel_map.py             ← CCCL 算法 → vllm kernel 映射
-├── muh_dispatch.py               ← 运行时 policy 分派
-├── vllm/                         ← vllm 引擎源码 (11MB Python)
-├── vllm_adapter/                 ← Qwen3.5 模型适配 + 部署脚本
-├── qwen3_6_scripts/              ← Qwen3.6 patch 集合 (576KB, 25+ patches)
-├── prefix_prefill.py             ← Triton prefill kernel (可调优)
-├── paged_attn.py                 ← Triton decode kernel (可调优)
-├── attention.py                  ← Attention 实现
-└── enginex-vllm-bi100-qwen36-main.zip ← 竞赛基础引擎 (97MB)
-```

GROUND_TRUTH_STATUS_v2.md

@@ -1,101 +0,0 @@
-# project_6 真实状态 v2
-
-更新时间: 2026-08-06, 基于完整代码阅读
-
-## 核心事实
-
-**enginex 没有 .cu 源码。gen_patch 的 C++ injection 全部失效。** 但这不是终点。
-
-实际可优化的三条路径：
-
-### 路径 1: Triton kernel 参数调优 (直接有效)
-
-文件: `prefix_prefill.py` (895行), `paged_attn.py` (794行)
-状态: 22 个 flash_attn 配置 + 9 个 prefill 配置已计算 SMEM，未上机实测
-关键参数:
-- prefill: BLOCK_M, BLOCK_N, NUM_WARPS (已有 SMEM 约束扫描)
-- decode: _PARTITION_SIZE=512 (硬编码), V1/V2 切换阈值
-- 竞赛权重: Output TPS×16.796(83%) + Input TPS×2.799(14%)
-
-gen_patch.py 第 87-103 行已经指向了这些真正的 injection points:
-```python
-('prefill', 'BLOCK_M'): [('prefix_prefill.py', 'BLOCK')],
-('flash_attn', 'BLOCK_M'): [('vllm/attention/ops/triton_flash_attention.py', 'BLOCK_M')],
-('moe', 'BLOCK_SIZE_M'): [('vllm/model_executor/layers/fused_moe/fused_moe.py', 'BLOCK_SIZE_M')],
-```
-
-### 路径 2: 模型适配 (功能门控)
-
-文件: `vllm_adapter/qwen3_5.py` (588行), `qwen3_6_scripts/` (25+ patches)
-状态: MoE 256 experts top-8 注册完成，treat ALL layers as full attention
-待验证: TP=4 加载, reasoning 分离, tool_call parsing
-竞赛门控: 50+ 功能测试全通过 + 效果偏差 ≤±4%
-
-### 路径 3: vllm Python 层配置优化 (低风险高收益)
-
-文件: `computility-run.yaml`, `baseline.muh`
-关键发现 from paged_attn.py:
-- 第 99 行: `use_v1 = True` 硬编码禁用了 V2 — 对 100K token 序列这是性能杀手
-- `_PARTITION_SIZE = 512` 硬编码 — 应该根据 SM count=16 动态调整
-- `max_num_seqs: 1` — 限制了批处理并行度
-- `--enable-prefix-caching` — 已开启，但 cache copy kernel 未优化
-
-## CCCL 资产的真实价值
-
-CCCL 的价值不在于 C++ 注入（已证实失效），而在于：
-
-1. **参数空间知识**: 27 个 tuning_*.cuh 告诉我们 NVIDIA 在 3 代 GPU 上搜索了哪些参数维度
-   - reduce: ipt×tpb×ipv = 1044 个组合
-   - scan: ipt×tpb×ns×dcid×l2w×trp×ld = ~26B 个（剪枝后可管理）
-   - 这些维度完全适用于 Triton kernel 的等价参数
-
-2. **benchmark 数据**: 199 条标注告诉我们在不同 problem size 下的加速比分布
-   - 小数据量(<16M): 大多数优化无效（speedup≈1.0）
-   - 大数据量(>256M): 加速比显著（最高 1.58x）
-   - 这意味着 decode（小 batch）和 prefill（大 batch）需要不同策略
-
-3. **约束模型**: scale_mem_bound, SMEM 公式, occupancy 计算
-   - BI-V100: 16 SM, 48KB SMEM, 900GB/s BW
-   - per-SM BW = 56 GB/s ≈ B200 水平
-   - bytes_in_flight = 64KB (bench_bi100.py 已验证)
-
-4. **算法映射**: muh_kernel_map.py 的 VLLM_KERNEL_MAP 精确映射了每个 vllm kernel 对应的 CCCL 算法
-   - paged_attention → reduce (summary_statistics.cu Welford pattern)
-   - softmax → scan
-   - sampling → topk + radix_sort
-   - normalization → transform + reduce
-
-## bench_bi100.py 的实际作用
-
-bench_bi100.py (713行) 是真正的工具 — 它用 PyTorch CUDA 操作模拟 CCCL benchmark:
-- 不需要编译 C++，不需要 nvbench
-- 直接在 BI-V100 上跑 torch.sum/torch.cumsum/torch.topk
-- 输出 CCCL 格式: `ipt_N.tpb_M.ipv_K speedup0 speedup1 speedup2 speedup3`
-- 搜索空间定义完整: reduce 1044 组合, scan 剪枝后可管理, topk/transform 都有
-
-**但它需要 BI-V100 硬件才能跑。** 在 Phanthy Cloud 上部署就能开始标定。
-
-## 代码覆盖率 (muh vs CCCL)
-
-| 算法 | muh 行 | CCCL 行 | 比率 | 竞赛价值 |
-|------|--------|---------|------|---------|
-| reduce | 297 | 478 | 62% | 最高 — Output TPS 83% |
-| topk | 113 | 121 | 93% | 高 — 每次 decode |
-| scan | 370 | 1525 | 24% | 高 — softmax |
-| transform | 185 | 549 | 34% | 中 — RMSNorm/SiLU |
-| select_if | 459 | 2729 | 17% | 中 — token filter |
-| radix_sort | 222 | 2381 | 9% | 中 — full sort |
-| scan_by_key | 145 | 2008 | 7% | 中 — per-seq scan |
-| reduce_by_key | 171 | 1735 | 10% | 中 — score aggregation |
-| unique_by_key | 166 | 1539 | 11% | 低 — KV dedup |
-| 其余 18 个 | 33-189 | 78-788 | varies | 低 |
-
-muh 总计 3618 行 / CCCL 17000+ 行 = 21% 平均覆盖率。
-reduce 和 topk 覆盖率最高（62%、93%），正好是竞赛权重最大的两个算法。
-
-## 下一步具体行动
-
-1. **在 Phanthy Cloud 上跑 bench_bi100.py** — 产出 BI-V100 真实 benchmark 数据
-2. **把 benchmark 结果回填到 Triton kernel 参数** — prefix_prefill.py 的 BLOCK/NUM_WARPS
-3. **修复 paged_attn.py 的 V2 禁用** — 对长序列性能至关重要
-4. **功能测试回归** — 确保 qwen3_5.py 适配通过 50+ 用例

MUH_PROJECT_CHECKPOINT.md

@@ -1,208 +0,0 @@
-# MUH Project Checkpoint
-
-> **最后更新**: 2026-07-30
-> **GitHub Project**: github.com/users/dylanyunlon/projects/6
-> **代码仓库**: github.com/dylanyunlon/project_6
-> **竞赛截止**: 2026-09-30
-
----
-
-## 一、项目是什么
-
-参加信创模盒 ModelHub XC 的"模型适配引擎竞赛-第一届"。目标是优化 vllm 引擎，让 Qwen3.6-35B-A3B 在天数智芯天垓100（4×BI-V100 GPU）上跑出最高的 Token 吞吐加权值。
-
-**计分公式**:
-```
-Token吞吐加权值 = Output TPS × 16.796 + Input TPS × 2.799 + Cache TPS × 0.56
-```
-
-Output TPS 权重占 83%——decode 阶段优化收益最大。
-
-**奖项**:
-- 基础奖 200,000 积分（1:1 兑现金）: 通过全部功能/效果测试 + 性能达标(≥8000)
-- 高级奖 +100,000: 加权值提升 ≥ 30%
-- 特级奖 +50,000: 加权值提升 ≥ 50%
-
-## 二、竞赛测评流程
-
-参赛者提交的是 **Git 仓库地址**（在 dev.modelhub.org.cn 上）。平台自动执行：
-
-1. **构建镜像**: 读取仓库根目录的 `Dockerfile`，基于基础镜像 `harbor.4pd.io/modelhubxc/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3` 构建
-2. **启动服务**: 读取 `computility-run.yaml` 的 `command`，在 4×天垓100 容器里启动 vllm api server（模型权重平台预挂载在 `/model`）
-3. **功能测试（门控）**: 50+ 个 OpenAI 兼容 API 测试用例，全部通过才进入下一步
-4. **效果测试（门控）**: 标准 benchmark 偏差 ≤ ±4%
-5. **性能测试（排名）**: 计算加权值
-
-**你能改的**: Dockerfile + vllm 源码 + computility-run.yaml 启动参数。模型本身不能改。
-
-## 三、muh 是什么
-
-muh 是我们设计的 **tuning DSL（领域特定语言）**，用于：
-
-1. 把 CCCL 的 tuning pattern（block_threads / items_per_thread / load_algorithm / cache_modifier 等）抽象成硬件无关的参数空间
-2. 针对天垓100 的硬件特性搜索最优参数组合
-3. Codegen 输出实际的 vllm kernel 修改 + computility-run.yaml + Dockerfile
-
-**为什么需要它**: CCCL 有 27 个 tuning_*.cuh 文件（17000+ 行），每个算法都有针对不同 NVIDIA SM 架构的特化参数。天垓100 不是 NVIDIA GPU，不能直接用这些参数，但 tuning 的维度（block size、warp 策略、shared memory 用量、prefetch 策略）是通用的。muh 让迁移过程变成"改配置 + 跑 benchmark"而不是"手改 kernel + 祈祷"。
-
-**muh 的状态**: v0.3 — 6个算法的C++ tuning headers已就绪(reduce/scan/topk/transform/batch_memcpy/for)，compile_test 33项通过，gen_patch.py从C++ headers提取bi100值生成vllm patches。参数值从CCCL SM100复制，等BI-V100实测替换。
-
-## 四、已完成的工作
-
-### 4.1 Project 6 已有 16 个真实 GitHub Issue（不是 Draft）
-
-都在 `dylanyunlon/project_6` 仓库里，已关联到 GitHub Project 6，有 label 和 Priority：
-
-| # | 标题 | Labels | Priority |
-|---|------|--------|----------|
-| 1 | [FEA] 非流式基础对话 | 基本功能,vllm,天垓100,Qwen3.6 | P0 |
-| 2 | [FEA] 流式对话 SSE | 基本功能,vllm | P0 |
-| 3 | [FEA] Tool Calling | 基本功能,vllm,Qwen3.6 | P0 |
-| 4 | [FEA] Reasoning/Thinking 分离 | 基本功能,thinking,Qwen3.6 | P0 |
-| 5 | [FEA] Prefix Cache | 基本功能,性能测试,vllm | P0 |
-| 6 | [FEA] 采样参数边界 | 采样参数,vllm | P1 |
-| 7 | [FEA] max_tokens 边界 | max_tokens,vllm | P1 |
-| 8 | [FEA] 结构化输出 | 结构化输出,vllm | P0 |
-| 9 | [FEA] 多语言 Emoji | 多语言,Qwen3.6 | P1 |
-| 10 | [FEA] 多模态 base64 PNG | 多模态,基本功能,Qwen3.6 | P0 |
-| 11 | [FEA] 参数校验 | 参数校验,vllm | P1 |
-| 12 | [FEA] 基础能力 | 基础能力,vllm,Qwen3.6 | P0 |
-| 13 | [FEA] 输出截断 | 截断测试,vllm | P1 |
-| 14 | [FEA] 效果测试 | 效果测试,Qwen3.6,天垓100 | P0 |
-| 15 | [EPIC] 性能基准 | 性能测试,天垓100,vllm | P0 |
-| 16 | [EPIC] 开发环境与代码提交 | infra,天垓100 | P1 |
-
-这 16 个覆盖了竞赛功能测试的所有 50+ 用例。每个 issue 的 body 里都有 PND 级别的测试用例表（前置条件 + 原子步骤 + 二值判定标准）。
-
-### 4.2 仓库里已有 NVIDIA CCCL 代码
-
-`project_6/cccl_upstream/` 目录下包含完整的 CCCL：
-- `cub/` — GPU 原语（reduce, scan, sort, topk, block/warp/device 三层）
-- `thrust/` — 高层算法 + 60 个示例
-- `libcudacxx/` — CUDA C++ 标准库
-- `cudax/` — 实验性功能（allocators, memory resources）
-- `cub/cub/device/dispatch/tuning/` — 27 个硬件特化 tuning 文件（17000+ 行）
-
-### 4.3 Label 体系已建立
-
-仓库上已创建 16 个 label：基本功能、thinking、采样参数、max_tokens、基础能力、结构化输出、多语言、多模态、参数校验、截断测试、效果测试、性能测试、infra、vllm、天垓100、Qwen3.6
-
-### 4.4 Project 6 里有 15 个遗留 Draft Issue 需要清理
-
-这些是早期用 addProjectV2DraftIssue 创建的，没有 repo 关联、没有 label。应该从 Project 面板里手动删除。
-
-## 五、还没做的（下一步）
-
-1. ~~muh 语言 PRD 设计~~ ✅ Done — muh是C++ header-only lib，不是独立语言
-2. ~~从 CCCL tuning_*.cuh 提取参数空间~~ ✅ Done — 6个算法的policy_selector已实现
-3. **在BI-V100上跑benchmark** — 用实测数据替换bi100_*中的SM100复制值
-4. **获取 enginex-vllm-bi100-qwen36 的实际代码** — 需要在 Phanthy Cloud 开发环境里操作
-5. **设计 muh → vllm kernel 的 codegen 管道**
-6. **实际在天垓100 上跑 benchmark**
-
-## 六、参考项目
-
-- **NVIDIA CCCL Project #6**: github.com/orgs/NVIDIA/projects/6（1990 items，Issue-first 模式，label 做模块分类）
-- **pub/sub-loop Project #4**: github.com/users/dylanyunlon/projects/4（1632 items，Draft-first 模式，已验证 1111 个有真实测试步骤，154 个有"按AC验证"占位符）
-- **PND 测试库**: 818 条车载软件测试用例，作为 PRD 测试用例质量基准
-
-## 七、关键文件路径
-
-```
-project_6/
-├── cccl_upstream/           # NVIDIA CCCL 完整代码
-│   ├── cub/cub/device/dispatch/tuning/  # 27 个 tuning policy 文件
-│   ├── cub/cub/warp/        # warp-level 原语
-│   ├── cub/cub/block/       # block-level 原语
-│   ├── thrust/examples/     # 60 个优化模式示例
-│   └── cudax/...allocators/ # 内存分配器
-├── Dockerfile               # TODO: 待创建
-├── computility-run.yaml     # TODO: 待创建
-└── muh/                     # TODO: muh 语言实现
-```
-
-## 八、竞赛关键参数（来自 computility-run.yaml 参考）
-
-```yaml
-concurrency: 1
-command:
-  - python3 -m vllm.entrypoints.openai.api_server
-  - --model /model
-  - --served-model-name llm
-  - --max-model-len 100000
-  - --gpu-memory-utilization 0.9
-  - -tp 4
-  - --max-num-seqs 1
-  - --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
-```
-
-基础镜像: `harbor.4pd.io/modelhubxc/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3`
-
-## 九、CCCL Tuning 文件全量模型输入记录
-
-**所有 27 个 tuning_*.cuh 文件的完整源码已在本 context 中作为模型输入读取。** 关键发现：
-
-### policy_selector 统一模式
-
-每个算法都有一个 `policy_selector` struct，接受 `::cuda::compute_capability cc` 参数，内部按 SM 版本做 if-else 分支：
-
-```
-if (cc >= {10, 0}) → sm100 tuning (Blackwell)
-if (cc >= {9, 0})  → sm90 tuning (Hopper)  
-if (cc >= {8, 0})  → sm80 tuning (Ampere)
-if (cc >= {7, 0})  → sm70 tuning (Volta)
-if (cc >= {6, 0})  → sm60 tuning (Pascal)
-fallback           → sm50 tuning
-```
-
-**muh 的核心工作就是给每个 policy_selector 添加一个 `cc == {iluvatar, 100}` 分支，填入在天垓100 上跑出的最优 benchmark 数据。**
-
-### 各算法提取的参数维度
-
-| 算法 | 文件 | 行数 | 参数维度 |
-|------|------|------|---------|
-| reduce | tuning_reduce.cuh | 478 | threads, items, vec_size, reduce_algorithm, load_modifier, determinism |
-| scan | tuning_scan.cuh | 1525 | threads, items, load_algo, load_mod, store_algo, scan_algo, delay_policy + lookahead variant |
-| radix_sort | tuning_radix_sort.cuh | 2381 | histogram(threads,items,partitions,radix_bits) + exclusive_sum + onesweep(threads,items,store,rank,scan,partitions,radix_bits) + downsweep + upsweep + single_tile |
-| reduce_by_key | tuning_reduce_by_key.cuh | 1735 | threads, items, load_algo, load_mod, scan_algo, delay_policy |
-| select_if | tuning_select_if.cuh | 2729 | threads, items, load_algo, load_mod, scan_algo, delay_policy |
-| histogram | tuning_histogram.cuh | 363 | threads, pixels_per_thread, vec_size, load_algo, load_mod, rle_compress, mem_preference, work_stealing |
-| topk | tuning_topk.cuh | 121 | threads, items (simple, no SM-specific tuning yet) |
-| batched_topk | tuning_batched_topk.cuh | 186 | worker_policy array × 6 tiers + multi_worker_policy |
-| merge | tuning_merge.cuh | 180 | threads, items, load_mod, store_algo, bulk_copy_keys, bulk_copy_values |
-| merge_sort | tuning_merge_sort.cuh | 193 | threads, items, load_algo, load_mod, store_algo |
-| transform | tuning_transform.cuh | 549 | threads, items, load_algo, store_algo, load_mod |
-| rle_encode | tuning_rle_encode.cuh | 626 | threads, items, load_algo, load_mod, scan_algo, delay_policy |
-| rle_non_trivial | tuning_rle_non_trivial_runs.cuh | 691 | threads, items, load_algo, load_mod, store_time_slicing, scan_algo, delay |
-| adjacent_diff | tuning_adjacent_difference.cuh | 118 | threads, items, load_algo, load_mod, store_algo (single policy, no SM branching) |
-| for | tuning_for.cuh | 78 | threads, items (trivial, 256×2) |
-| find | tuning_find.cuh | 90 | threads, items, vec_size, load_mod |
-| batch_memcpy | tuning_batch_memcpy.cuh | 227 | small_buffer + large_buffer sub-policies |
-| scan_by_key | tuning_scan_by_key.cuh | ~2000 | same as reduce_by_key pattern |
-| unique_by_key | tuning_unique_by_key.cuh | ~1500 | same pattern |
-| three_way_partition | tuning_three_way_partition.cuh | ~780 | same pattern |
-| segmented_* | 4 files | ~1300 total | segmented variants of reduce/scan/sort |
-
-### Benchmark 注释格式
-
-每个 sm100 tuning 都有注释格式：
-```
-// ipt_22.tpb_384.ns_1904.dcid_6.l2w_830.trp_1.ld_0 1.148442 0.997167 1.139902 1.462651
-```
-- `ipt` = items_per_thread
-- `tpb` = threads_per_block  
-- `ns` = delay nanoseconds
-- `dcid` = delay constructor ID
-- `l2w` = L2 cache window
-- `trp` = transpose (0=DIRECT, 1=WARP_TRANSPOSE)
-- `ld` = load modifier (0=DEFAULT, 1=LDG, 2=CA)
-- 4 个数字 = 4 种 problem size 下的加速比 (vs 前代 SM)

MUH_TUNING_GAP_ANALYSIS.md

@@ -1,85 +0,0 @@
-# muh Tuning Gap Analysis — CCCL vs BI-V100 适配
-## 2026-08-07
-
-### 方法论
-
-直接读取 CCCL 源码（26 个 tuning_*.cuh），提取竞赛相关的 benchmark annotations，
-对比 muh 已有的 BI-V100 struct 值。每个算法的优先级由竞赛评分公式决定：
-
-```
-Score = Output_TPS × 16.796 + Input_TPS × 2.799 + Cache_TPS × 0.56
-```
-
-Output TPS = 83%, Input TPS = 14%, Cache TPS = 3%
-
----
-
-### P0: 直接影响竞赛评分的算法
-
-#### 1. REDUCE (Output TPS 83%) — ★★★★★
-- **竞赛路径**: paged_attention score reduction, float32, plus
-- **CCCL SM100**: `ipt_16.tpb_512.ipv_2 → 1.061/1.000/1.065/1.167`
-- **muh BI-V100**: `bi100_plus_float32_o4 {512, 24, 2}` — tile=12288 (1.5× SM100)
-- **状态**: ✅ 完成 (62% 行覆盖)
-- **待定**: SM=16 items 适配 (P0 BUG)、LOAD_LDG vs LOAD_DEFAULT benchmark
-
-#### 2. SCAN (Output TPS 83%) — ★★★★☆
-- **竞赛路径**: softmax denominator prefix sum, float32, plus
-- **CCCL SM100**: `ipt_22.tpb_384.ns_1904.dcid_6.l2w_830 → 1.148/0.997/1.140/1.463`
-- **muh BI-V100**: `bi100_lookback_4B_o4 {384, 22}` — 与 SM100 同 tile
-- **状态**: ✅ 核心完成 (39% 行覆盖，lookback + SM90 fallback)
-- **待定**: Lookback delay 参数需实测校准、8B structs 99% SMEM 需验证
-
-#### 3. TRANSFORM (Input TPS 14% + all activations) — ★★★★☆
-- **竞赛路径**: SiLU/GeLU/RMSNorm, bfloat16
-- **CCCL**: bytes_in_flight 是核心参数, B200=64KB, H100=48KB
-- **muh BI-V100**: bytes_in_flight=64KB (confirmed by babelstream bench)
-- **状态**: ✅ 核心完成
-- **待定**: Vectorized vs prefetch algorithm 选择需实测
-
----
-
-### P1: 间接影响性能的算法
-
-#### 4. TOPK (sampling, Output TPS) — ★★★☆☆
-- **竞赛路径**: logit sampling, float32 keys
-- **CCCL**: bits_per_pass, thread count, BLOCK_SCAN_WARP_SCANS
-- **muh BI-V100**: 有 inline tuning (threads=512, bits_per_pass=11)
-- **状态**: ✅ 基本完成
-- **待定**: Onesweep vs multi-sweep 选择
-
-#### 5. SELECT_IF (MoE routing) — ★★☆☆☆
-- **竞赛路径**: expert selection, float32, not_flagged, no_rejects, offset_4
-- **CCCL SM80**: `{threads=256, items=18, WARP_TRANSPOSE, no_delay=1130}`
-- **muh BI-V100**: 零 bi100 structs, 用 get_sm100_adapted() inline 计算
-- **状态**: ⚠️ 只需 1/77 个 specialization, 但完全缺失
-- **待定**: 需添加 bi100_select_float32_nf_nr_o4 struct
-
-#### 6. RADIX_SORT (topk helper) — ★★☆☆☆
-- **竞赛路径**: float32 key sort for sampling
-- **CCCL**: 2381 行, onesweep + histogram, SM100 有复杂分支
-- **muh BI-V100**: 222 行 (9% 覆盖)
-- **状态**: ⚠️ 需要 onesweep 路径
-- **待定**: bits_per_pass 和 histogram SMEM
-
----
-
-### P2: 理论覆盖但不直接影响评分
-
-| 算法 | CCCL 行数 | muh 行数 | 覆盖率 | 竞赛影响 |
-|------|----------|---------|-------|---------|
-| reduce_by_key | 1735 | 217 | 13% | 低 |
-| scan_by_key | 2008 | 161 | 8% | 低 |
-| unique_by_key | 1510 | 179 | 12% | 低 |
-| three_way_partition | 708 | 67 | 9% | 低 |
-| segmented_reduce | 471 | 112 | 24% | 低 |
-| 其余 14 个 | ~4000 | ~800 | ~20% | 无 |
-
----
-
-### 关键差距总结
-
-1. **gen_patch.py 管道断裂** — 产出零 patch。已被 gen_config.py 替代。
-2. **muh headers 20% 完成** — 但竞赛相关的 5 个算法 (reduce/scan/transform/topk/select_if) 核心参数已就位。
-3. **缺 benchmark 验证** — 所有 BI-V100 speedup 标 TBD，需要在 Phanthy Cloud 上跑。
-4. **Python layer 是真正的注入点** — 已在 triton_flash_attention.py 添加 8 个 BI-V100 configs, prefix_prefill.py 修 BLOCK=64, _custom_ops.py 修 SMEM=48KB。gen_config.py 又发现 19 个新候选 configs。

PIPELINE_GROUND_TRUTH.md

@@ -1,50 +0,0 @@
-# muh Pipeline Ground Truth — 2026-08-07
-
-## 管道实际状态（不是设计稿，是已部署代码的真实描述）
-
-### scale_mem_bound: FULL PARITY ✓
-11/11测试用例与CCCL `cub::detail::scale_mem_bound` 完全匹配。
-返回值顺序 `{items_per_thread, threads_per_block}` — items-first，与CCCL一致。
-
-### C++ Tuning Headers: 27/27 ✓
-所有26个算法（+common）都有bi100 header，`policy_selector::operator()` 接受
-`hardware_capability` 参数。SMEM overflow保护覆盖所有type_size。
-
-### Injection现状（enginex没有.cu源码）
-
-| 注入位置 | 状态 | 值 | commit |
-|---------|------|-----|--------|
-| prefix_prefill.py BLOCK | ✓ 已手动修改 | BLOCK=64, WARPS=4 | 多个commit |
-| paged_attn.py _PARTITION_SIZE | ✓ 保持默认 | 512 | — |
-| paged_attn.py V1/V2 dispatch | ✓ 已手动修改 | use_v1 threshold | cbd1f08 |
-| _custom_ops.py SMEM | ✓ 已手动修改 | 48KB | 16f0b30 |
-| triton_flash_attention.py | ✓ 已添加BI-V100 configs | BLOCK=32/64 | 多个commit |
-| protocol.py 兼容性 | ✓ 已修复 | max_completion_tokens等 | 2c353da |
-
-### gen_patch.py 角色
-设计时期望: C++ header → unified diff → vllm .cu文件
-实际情况: enginex只有Python + .so, 没有.cu源码
-当前角色: 文档工具 + 验证（确认header值与已部署Python代码一致）
-
-### CCCL SM100 Benchmark数据（从源码提取，已存入cccl_sm100_benchmark_values.json）
-
-**Reduce** (paged_attention score reduction, Output TPS 83%权重):
-- float32+plus: items=16, threads=512, vec=2, speedup=[1.061, 1.000, 1.065, 1.167]
-- float64+plus: items=16, threads=640, vec=1, speedup=[1.018, 1.000, 1.016, 1.057]
-
-**Scan** (softmax prefix-sum):
-- 4B lookback: items=22, threads=384, delay=1904ns/dcid=6/l2w=830, speedup=[1.148, 0.997, 1.140, 1.463]
-- 8B lookback: items=23, threads=416, delay=772ns/dcid=5/l2w=710, speedup=[1.089, 1.016, 1.086, 1.265]
-
-**muh BI-V100适配**:
-- reduce float32: items=24(+50%), threads=512(=), vec=2(=) → 补偿16 SMs
-- scan 4B: 通过scale_mem_bound自动适配(items=22 @4B安全, @8B降级到16)
-- delay参数: ns×0.5, l2w×0.6 (启发式, 待实测)
-
-### 竞赛门槛
-- 功能测试: 50+ TC, 项目看板14个FEA item覆盖
-- 效果测试: benchmark偏差 ≤ ±4%
-- 性能测试: Token吞吐加权值 ≥ 8000
-  - Output TPS × 16.796 (83%) → reduce/scan/topk
-  - Input TPS × 2.799 (14%) → scan/transform
-  - Cache TPS × 0.56 (3%) → batch_memcpy

PIPELINE_REALITY_CHECK.md

@@ -1,71 +0,0 @@
-# muh 管道现实检查 — 2026-08-07
-
-## 核心发现
-
-### 1. gen_patch.py 输出为零
-
-```
-$ python3 muh/gen_patch.py --dry-run
-READ reduce: bi100_plus_float32_o4 → {items: 24, threads: 512, vec: 2}
-READ scan: bi100_sm90_float32 → {threads: 128, items: 24}
-...
-No patches generated.
-```
-
-原因: `VLLM_INJECTION_POINTS` 的 key `('reduce', 'partition_size')` 和 struct 提取出的 field `items`/`threads`/`vec` 不匹配。gen_patch 的"读"和"写"两端从未对齐。
-
-### 2. 注入目标是 Python 不是 C++
-
-enginex-vllm-bi100 **没有 `.cu` 源码**。所有 CUDA kernel 是预编译的 ixformer `.so`。
-
-实际可调的全部是 Python 层:
-
-| 文件 | 可调参数 | 竞赛影响 |
-|------|---------|---------|
-| `paged_attn.py` | `_PARTITION_SIZE=512`, V1/V2 dispatch logic | Output TPS (83%) |
-| `prefix_prefill.py` | `BLOCK=64`, `BLOCK_N=64`, `NUM_WARPS=4` | Input TPS (14%) |
-| `vllm/attention/ops/triton_flash_attention.py` | 17 个 autotune configs | Prefill throughput |
-| `vllm/_custom_ops.py` | `return 49152` (SMEM fix) | 所有 Triton kernels |
-| `computility-run.yaml` | `--max-num-seqs`, `--gpu-memory-utilization` | 调度效率 |
-
-gen_patch.py 中的 `csrc/*.cu` 注入点全部是 dead code (注释已标注)。
-
-### 3. muh C++ headers 的实际价值
-
-muh 的 26 个 tuning headers 和 `scale_mem_bound` 实现是正确的理论分析工具。它们的价值不在于直接注入 vllm，而在于:
-
-- 推导 SMEM 约束 (Triton `BLOCK_M × head_dim × elem_size` 上限)
-- 推导 occupancy 模型 (BI-V100 16 SMs 的 wave efficiency)
-- 推导 bytes_in_flight (56 GB/s per-SM → 64KB prefetch window → `num_stages=2`)
-- 为 CCCL benchmark 验证提供 ground truth
-
-这些推导已经手工应用到了 Python 代码中:
-- `triton_flash_attention.py` 的 8 个 BI-V100 configs 引用了 CCCL babelstream/scan 分析
-- `prefix_prefill.py` 的 BLOCK_N=64 推导基于 48KB SMEM 约束
-- `_custom_ops.py` 的 49152 来自 hardware.cuh
-
-### 4. 管道闭环的正确路径
-
-```
-CCCL tuning analysis    Python layer injection   Triton autotune
-(理论推导)               (参数修改)                (运行时选择)
-        │                       │                       │
-        ▼                       ▼                       ▼
-muh headers           paged_attn.py             triton.Config([...])
-common.cuh            prefix_prefill.py         autotune picks best
-hardware.cuh          _custom_ops.py            at runtime
-        │                       │                       │
-        └───────────────────────┴───────────────────────┘
-                                │
-                        竞赛评测得分
-```
-
-不是: `muh headers → gen_patch → #define injection → recompile`
-而是: `muh analysis → Python config → Triton autotune → runtime perf`
-
-## 下一步
-
-1. 删除 gen_patch.py 中所有 dead `csrc/*.cu` 注入点
-2. 重写 gen_patch 为 `gen_config.py`: 从 muh headers 推导 → 直接输出 Python patch
-3. 用 CCCL benchmarks 验证: reduce/sum.cu, scan/exclusive/sum.cu, topk/keys.cu
-4. 扩展 triton_flash_attention.py autotune 搜索空间 (当前 17 configs, 可加到 30+)

PIPELINE_STATUS.md

@@ -1,86 +0,0 @@
-# muh Pipeline Status — Ground Truth
-
-**Last verified**: 2026-08-07T01:45:31Z by automated analysis
-
-## Architecture Summary
-
-```
-CCCL policy_selector(compute_capability) → ReducePolicy{threads, items, vec, algo, load_mod}
-     ↕ mirrors
-muh  policy_selector(hardware_capability) → same struct types, BI-V100 values
-     ↕ gen_patch.py extracts bi100_* values
-vllm patch_ops.sh → full-file Python replacements with tuning values baked in
-```
-
-## Injection Reality
-
-### What gen_patch.py THINKS (csrc/*.cu — DEAD)
-```
-tuning_reduce.cuh → csrc/attention/attention_kernels.cu NUM_THREADS   ← NO .cu SOURCE
-tuning_scan.cuh   → csrc/attention/paged_attention_v1.cu SCAN_BLOCK_SIZE ← NO .cu SOURCE
-tuning_topk.cuh   → csrc/sampling/sampling_kernels.cu SAMPLING_BLOCK_SIZE ← NO .cu SOURCE
-```
-
-### What ACTUALLY happens (Python runtime — ALIVE)
-```
-_custom_ops.py    → SMEM 49152 (was 32768)     ← DEPLOYED ✓
-paged_attn.py     → _PARTITION_SIZE=512         ← DEPLOYED ✓ (V2 partition, NOT CTA tile)
-xformers.py       → _Q_CHUNK=256, sdpa_fallback ← DEPLOYED ✓
-sampler.py        → torch.topk fast path        ← DEPLOYED ✓
-prefix_prefill.py → Triton BLOCK_M/N/warps      ← DEPLOYED ✓ (but Triton not available)
-computility-run.yaml → vllm server args         ← DEPLOYED ✓
-```
-
-### The Gap
-muh C++ headers define precise per-type-per-op tuning values (14 reduce structs, 22 scan structs).
-But the vllm engine on BI-V100 runs ixformer .so (precompiled, not tunable) + Python fallbacks.
-The C++ headers' values cannot be injected into the precompiled .so.
-They CAN inform:
-1. Python fallback implementations (paged_attn.py, xformers.py) — tile sizes, chunk sizes
-2. Triton JIT configs — if Triton were available (it's not on BI-V100 base image)
-3. Future EngineX releases that expose tuning knobs
-
-## Asset Inventory
-
-| Asset | Count | Status |
-|-------|-------|--------|
-| CCCL tuning headers (upstream) | 27 | Complete |
-| muh BI-V100 headers | 27 | Complete (14 reduce + 22 scan + others) |
-| muh schema YAMLs | 27 | Complete |
-| CUB benchmarks | 91 | Synced to NVIDIA/cccl main |
-| CUB tests | 243 | Complete |
-| CUB examples | 18 | Complete |
-| Thrust examples | 60 | Complete |
-| Deployed patches | 15 files | Via patch_ops.sh full replacement |
-| bench_bi100.py search spaces | 5 algos | Defined, needs BI-V100 hardware to run |
-
-## Tool Chain Status
-
-| Tool | Input | Output | Status |
-|------|-------|--------|--------|
-| parse.py | baseline.muh | JSON config | ✓ Working |
-| gen_patch.py | tuning_*.cuh | Patch report | ⚠ Reports structs but generates 0 patches (injection mapping mismatch) |
-| gen_yaml.py | baseline.muh | computility-run.yaml | ✓ Working |
-| bench_bi100.py | algo+dtype | CCCL-format speedup data | Needs BI-V100 hardware |
-| patch_ops.sh | qwen3_6_scripts/ | Docker vllm patches | ✓ Working |
-| muh_dispatch.py | hw+dtype+head_dim | AttentionConfig | ✓ Working (needs torch) |
-| scale_mem_bound | (threads, items, type_size) | (items, threads) | ✓ CCCL parity verified |
-
-## Critical Numbers
-
-| Metric | Competition Threshold | Current Status |
-|--------|----------------------|----------------|
-| Functional tests | 50+ pass | 13 items In Progress (all FEA) |
-| Effect deviation | ≤ ±4% | Untested (needs hardware) |
-| Token throughput weighted | ≥ 8000 | Untested |
-| Output TPS weight | 83% (×16.796) | Reduce/scan/topk optimization focus |
-| SMEM limit | 49152 bytes | All 36 scan+reduce structs verified ✓ |
-| SM count | 16 (confirmed) | All headers updated |
-
-## Next Actions (Ranked by Competition Impact)
-
-1. **Run bench_bi100.py on BI-V100** → get real speedup data for reduce/scan/topk
-2. **Backfill speedup data to muh headers** → replace TBD/theoretical values
-3. **Optimize Python fallback tile sizes** → paged_attn.py, xformers.py Q_CHUNK
-4. **Tune computility-run.yaml** → max-num-seqs, max-batched-tokens, gpu-mem-util
-5. **Enable prefix caching benchmark** → cached_tokens > 0 for repeat prompts

PRD.md

@@ -9,152 +9,8 @@
 - 性能门槛 Token 吞吐加权值 ≥8000
 - Output TPS 权重占 83%（decode kernel 优化投入产出比最高）
 
-## 架构策略
-CCCL系统设计移植 + base引擎serving层改造。
+AllReduce 大概占 10ms。剩下的 36ms 是 Python dispatch。1400 次 PyTorch 函数调用 × 25 微秒。
 
-### 核心原则
-1. **不覆盖模型层代码** — Sub168证明base镜像CoreX原生代码能正确运行
-2. **只部署serving层** — patch_ops.sh控制部署范围
-3. **通过环境变量做硬件适配** — CCCL policy_selector模式
+这台机器有没有 NVLink 改变不了 Python 每次调用花 25 微秒的事实。NVIDIA 上用 CUDA Graph 一次性录制所有 kernel launch，replay 时零 Python 开销。但 BI-V100 CUDA 10.2 对 Graph 支持有限。
 
-### 部署文件清单（patch_ops.sh）
-- protocol.py — OpenAI API兼容层
-- serving_chat.py — 请求处理核心
-- qwen3coder_tool_parser.py — Qwen3 XML tool call解析
-- reasoning/ — thinking/reasoning分离
-- api_server.py — 入口点
-- chat_utils.py — 消息预处理
-- cli_args.py — 参数注册
-- registry.py — 仅当base缺少Qwen3_5时
-
-### 不部署的文件（base镜像原生）
-qwen3_5.py, model_runner.py, _custom_ops.py, sampler.py,
-scheduler.py, sequence.py, xformers.py, paged_attn.py,
-prefix_prefill.py, logits_processor.py, mamba_cache.py, arg_utils.py
-
-## Sub168参数基准（已对齐）
-- max_model_len=256000
-- max_num_seqs=2
-- gpu_memory_utilization=0.95
-- max_num_batched_tokens=4096
-- enable_chunked_prefill=True
-- enforce_eager=True
-- dtype=half
-- tensor_parallel_size=4
-
-## CCCL → base 映射记录
-| CCCL源码 | 映射到base位置 | 改动类型 |
-|----------|---------------|---------|
-| buddy_allocator.cu | computility-run.yaml env | PYTORCH_CUDA_ALLOC_CONF |
-| device_reduce policy_selector | computility-run.yaml params | 启动参数对齐Sub168 |
-| agent_reduce_by_key ConsumeTile | serving_chat.py | fast path/safe path分离 |
-| tuning_find_bound_sorted_values | yaml --dtype half | 类型大小自适应 |
-
-## 已修复的Sub508/509失败点
-1. ✅ n>1 OOM级联 → 允许n=2（匹配max_num_seqs=2）
-2. ✅ max_completion_tokens 400 → protocol.py接受
-3. ✅ tool_calls content=None → chat_utils.py容错
-4. ✅ d03 tool_call thinking耗尽 → 自动禁用thinking
-5. ✅ 内存碎片OOM → PYTORCH_CUDA_ALLOC_CONF
-6. ✅ 模型层代码破坏CoreX → patch_ops.sh只部署serving层
-
-## CCCL tuning_select_if.cuh → serving_chat.py 映射
-
-### 设计思想翻译
-CCCL三级分发：compute_capability → sm_tuning → benchmark参数
-我们三级分发：请求类型 → 处理路径 → Sub168实测参数
-
-### 参数对应关系
-| CCCL概念 | 我们的对应 |
-|---------|-----------|
-| compute_capability (SM80/90/100) | 请求类型 (tool_call/reasoning/basic) |
-| input_size (1/2/4/8 bytes) | 请求复杂度 (simple/multimodal/multi-turn) |
-| flagged/unflagged | has_tools/no_tools |
-| keep_rejects/discard | enable_thinking/disable_thinking |
-| threads_per_block | max_tokens cap |
-| items_per_thread | default_max_tokens计算 |
-| delay_constructor | token budget 分配策略 |
-| benchmark注释 (4个加速比) | Sub168日志实测数据 |
-
-### Sub168 benchmark数据（=我们的tuning表）
-| 请求类型 | 时间 | token数 | TPS |
-|---------|------|---------|-----|
-| d01 basic | 8.49s | 139 | 16.4 |
-| d03 tool_call | 2.12s | ~34 | ~16 |
-| d04 reasoning | 17.78s | 1192 | 67 |
-| d07 reasoning+content | 61.11s | 4451 | 72.8 |
-| replay avg | - | - | 11.86 |
-
-## CCCL agent_rle.cuh → streaming SSE 映射
-
-| CCCL agent_rle | serving_chat.py |
-|---------------|----------------|
-| streaming_context.num_uniques() | reasoning_token_counts[i] |
-| streaming_context.base_offset() | previous_num_tokens[i] |
-| BlockDiscontinuity (值变化检测) | reasoning_end_arr[i] (</think>检测) |
-| per-partition isolated state | per-choice state arrays |
-| ScatterDirect (压缩输出) | delta_message分发 |
-
-## CCCL adjacent_difference → streaming delta 映射
-
-| CCCL adjacent_diff | serving_chat.py |
-|-------------------|----------------|
-| SubtractLeftCopy | delta_text = output.text (保留原始+输出差值) |
-| previous element | previous_texts[i] |
-| current = prev + delta | current_text = previous_text + delta_text |
-| update prev = current | previous_texts[i] = current_text |
-
-## CCCL tuning_batched_topk.cuh → 采样策略映射
-
-### 设计思想
-6级worker_policy按tile size递减排列。运行时选最小够用的配置。
-multi_worker_policy用于超大segment的协作处理。
-
-### 映射
-| CCCL batched_topk | 我们的对应 |
-|-------------------|-----------|
-| worker_policy.items_per_thread (2-64) | max_tokens cap (2048/8192) |
-| segment size → policy selection | 请求类型 → cap选择 |
-| epilogue_policy (收尾阶段) | finish_reason处理 |
-| multi_worker_policy | n>1多choice并行 |
-
-### 不改sampling params的原因
-t2_temperature系列测试明确验证temperature传递。
-覆盖默认值会导致测试失败。当前策略正确。
-
-## CCCL block_reduce_warp_reductions → DeltaNet chunk_size
-
-### 设计思想
-Sequential path dominance → reduce per-iteration work.
-CCCL: thread count固定时，减少items_per_thread让每个thread做更少work。
-我们: Python loop iterations固定(=chunk_size)，减少chunk_size从32→16。
-
-## CCCL execution/exception.cuh → api_server.py _select_error_policy
-
-### 设计思想  
-Device code: exception_ptr永远false，不假装能恢复，直接fail fast。
-Host code: 用标准exception。
-我们: _select_error_policy按exception类型分发——OOM→503, dead→503, validation→400。
-
-## CCCL segmented_sort.cu → 完整base引擎迁移
-
-### 关键发现
-通过segmented_sort.cu的AST链追溯到base引擎zip包，发现我们缺少10个关键文件。
-
-### base引擎完整部署清单 (patch_ops.sh)
-
-| 文件 | 作用 | 缺失后果 |
-|-----|------|---------|
-| paged_attn.py | 绕过Triton hang | GPU永久挂起 |
-| patch_model_runner.py | 修prefix_cache_hit bug | chunked prefill第2+chunk crash |
-| mamba_cache.py | GatedDeltaNet状态管理 | 状态丢失→输出错误 |
-| sequence.py | 修completion_tokens膨胀 | 10K prompt×3 chunks = 30K虚假token |
-| scheduler.py | prefix cache metrics | 无cache统计 |
-| patch_xformers_sdpa_seq.py | head_dim=256 bypass | attention crash |
-| qwen3_5.py | 模型代码（条件部署） | ModuleNotFoundError |
-| serving层6文件 | API兼容 | 功能测试全失败 |
-| transformers==4.55.3 | Qwen3_5Config支持 | 配置加载失败 |
-
-### qwen3_5.py策略
-base原版1369行 → 无nan_to_num、无clamp、无_hw_policy。
-条件部署：如果Docker镜像已有>1000字节的qwen3_5.py就不覆盖。
+最大的问题是 太多小 kernel 走 Python dispatch。减少 launch 次数比优化任何单个 kernel 都有效。

PROJECT_SUMMARY.md

@@ -4,119 +4,3 @@
 天垓100 (BI-V100) 推理引擎竞赛，在 4×BI-V100 上运行 Qwen3.5-27B 推理服务。
 竞赛目标：Token吞吐加权值 ≥ 8000（Output TPS × 83% + Input TPS × 14% + Cache TPS × 3%）
 
-## 技术栈
-- Base image: bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
-- vLLM 0.6.3 (base) + serving层patch
-- ixformer (CoreX SDK, 含 flash_attn / paged_attention / silu_and_mul 等)
-- Tensor Parallel = 4, enforce_eager=True
-
-## 文件结构
-
-```
-project_6/
-├── PRD.md                    # 竞赛需求 + CCCL→base映射
-├── SYSTEM_DESIGN.md          # 架构设计: Docker/Build/Runtime/GDN dispatch
-├── Dockerfile                # Docker构建
-├── computility-run.yaml      # vLLM启动参数
-├── qwen3_6_scripts/          # serving层 + model patches (部署到vllm)
-│   ├── qwen3_5.py           (2040行) 模型代码: GDN + MoE + Attention
-│   ├── serving_chat.py       OpenAI API处理核心
-│   ├── protocol.py           请求/响应模型
-│   ├── api_server.py         FastAPI入口
-│   ├── patch_ops.sh          部署脚本 (全部patch的安装器)
-│   ├── flash_qla_sm70/       GDN CUDA kernel (gdn_forward.cu 1919行)
-│   └── ...                   其他patches
-├── ex_engine/                # EX引擎: 算法因子置换层
-│   ├── csrc/
-│   │   ├── ix_full_bridge.cpp    (331行) pybind11桥接→ixformer::infer 14个C++函数
-│   │   ├── ix_moe_bridge.cpp     (258行) MoE-only子集桥接
-│   │   └── moe_topk_softmax_v3.cu (148行) 独立CUDA topk kernel
-│   ├── python/
-│   │   ├── corex_moe.py      (196行) MoE分发: ix_bridge→ixformer::infer 7步pipeline
-│   │   ├── corex_gdn.py      (217行) GDN分发: chunked delta rule + decode
-│   │   ├── corex_fa2.py      (228行) FA2分发: packed/paged/chunked三模式
-│   │   ├── ix_bridge.py      (162行) ix_full_bridge.so加载器
-│   │   └── moe_topk.py        CUDA topk Python wrapper
-│   ├── build.sh              编译脚本 (corex clang/16)
-│   └── include/              C++ headers
-├── cccl_upstream/            (8900文件) NVIDIA CCCL strategic subset
-│   ├── cub/                  tuning headers + benchmarks + tests
-│   ├── thrust/               examples + tests
-│   └── libcudacxx/           C++ STL headers
-├── muh/                      muh工具链: BI-V100 tuning parameter生成
-│   ├── include/muh/tuning/   27个BI-V100 policy_selector headers
-│   └── gen_patch.py          C++ header → vllm unified diff
-├── upstream_ref/             上游参考代码
-│   ├── ds_vllm/              ds-vllm (vllm fork, 含topk_softmax_kernels.cu)
-│   └── xllm/                 xllm (ILU backend: kernels/ilu + layers/ilu)
-├── vllm/                     vllm源码副本 (参考用)
-└── docs/                     分析文档
-```
-
-## 关键文件说明
-
-### ex_engine/csrc/ix_full_bridge.cpp
-- `ix_topk_softmax()` → `ixformer::infer::topk_softmax`
-- `ix_moe_gen_idx()` → `ixformer::infer::moe_compute_token_index_api`
-- `ix_moe_expand_input()` → `ixformer::infer::moe_expand_input`
-- `ix_group_gemm()` → `ixformer::infer::moe_w16a16_group_gemm`
-- `ix_silu_and_mul()` → `ixformer::infer::silu_and_mul`
-- `ix_moe_combine_result()` → `ixformer::infer::moe_output_reduce_sum`
-- `ix_fused_moe_forward()` — 以上6步组合, 一次C++调用完成整个MoE
-- `ix_paged_attention()` → `ixformer::infer::xllm_paged_attention`
-- `ix_flash_attn_prefill()` → `ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables`
-- `ix_rms_norm()` / `ix_fused_add_rms_norm()` / `ix_rotary_embedding()` / `ix_reshape_and_cache()`
-
-### ex_engine/python/corex_moe.py
-- `moe_forward()` — 3级分发: ix_bridge全C++ → ix_bridge逐步 → Python loop
-- `topk_softmax()` — ix_bridge优先, fallback到Python softmax+topk
-- `moe_prefill()` / `moe_decode()` — 日志匹配comp 168格式
-
-### qwen3_6_scripts/qwen3_5.py
-- `GatedDeltaNet.forward()` — GDN层: corex_gdn dispatch
-- `Qwen3_5MoE.forward()` — MoE层: Tier 0-3分发 (ix_fused_moe → ix_bridge → corex_moe → PyTorch)
-
-## 当前状态
-- 370+ commits, 67 GitHub issues (63 open, 4 closed)
-- GitHub Project #6: 149 items (121 draft issues + 28 real issues)
-- CCCL upstream (5205 files) 作为工程基座, tuning/dispatch pattern 1:1映射
-- 真机 comp 168 日志已完整分析: 3个致命bug已定位并修复
-- 可提交竞赛平台测试
-
-## 本次任务完成内容
-comp 168 docker日志 + upstream_ref 系统设计分析 → 三个致命bug修复:
-
-1. **OOM修复**: computility-run.yaml max_model_len 256000→80000
-   - comp 168日志: `torch.cuda.OutOfMemoryError: Tried to allocate 32.00 MiB`
-   - 引擎OOM→崩溃→replay_tencent 881请求中704个 Connection refused
-   - BI-V100 KV cache容量~88112 blocks, 256000远超上限
-
-2. **topk_softmax ERROR日志消除**: _custom_ops.py silent fallback
-   - comp 168日志: `ixformer.functions has no attribute vllm_moe_topk_softmax` × 500+次
-   - 从 ixformer.h 确认 `ixformer::infer::topk_softmax` 在C++层存在但Python binding缺失
-   - 新代码: 尝试 ixformer._C.topk_softmax → 安静 PyTorch fallback
-
-3. **_custom_ops.py 部署**: patch_ops.sh 添加部署步骤
-   - 之前标记为 "DO NOT deploy", 现在修复后部署
-
-关键发现 (from upstream_ref/xllm):
-- xllm/core/kernels/ilu/ixformer.h: 完整的 ixformer::infer API (14函数)
-- xllm/core/layers/ilu/fused_moe.cpp: 生产级7步MoE pipeline (797行)
-- xllm/core/kernels/ilu/fused_moe.cpp: topk_softmax + gen_idx + expand + combine
-- 这些代码在 upstream_ref 中已存在, 接口与我们的 ix_full_bridge.cpp 完全一致
-
-## 历史任务摘要
-- comp 168 三个致命bug修复 (OOM + topk_softmax + _custom_ops部署)
-- corex_moe/corex_gdn/corex_fa2 dlopen模块重写 (ixformer::infer dispatch chain)
-- CCCL upstream导入(5205文件) + 27/27 muh tuning headers + CCCL→vllm pattern mapping
-- ix_full_bridge.cpp 14函数桥接 + moe_topk_softmax_v3.cu
-- GDN dtype guard + NaN clamp修复
-- serving层部署(protocol/serving_chat/api_server等) + Sub508/509功能修复
-- 67 GitHub issues + 121 draft issues + PRD/SYSTEM_DESIGN文档
-
-## 遗留问题/下次继续
-1. **GDN NaN (P0)** — prefill GDN 99.98% NaN, 替换为zeros=模型质量归零; 需要参考 xllm/npu_torch/qwen3_gated_delta_net_base.cpp 做 fp32 accumulation
-2. **真机编译ix_full_bridge.cpp** — JIT编译后MoE走Tier 0 (C++ 7步) 取代 Python loop
-3. **MoE性能** — 当前全走PyTorch for循环 (64 experts × 每token), Output TPS=11.86
-4. **121个draft issues→真issue** — GitHub API批量转换
-5. **提交竞赛平台** — 当前修复应能通过functional_acceptance基本测试, 不再OOM崩溃

SUB509_DEEP_DIAGNOSIS.md

@@ -1,141 +0,0 @@
-# Sub509 深度诊断 — 基于CCCL源码阅读的系统级分析
-
-## 一、Sub509 vs Sub168 关键数据对比
-
-| 测试 | 对手Sub168 | 我们Sub509 | 差距分析 |
-|------|-----------|-----------|---------|
-| d01_basic_nostream | 8.49s, content[11] tok=139 | 95.85s, content[0] reasoning[1102] tok=1085 | 11x慢; 我们产了1085个token全是reasoning |
-| d02_stream_usage | 2.75s, chunks=53 | 1.84s, chunks=9 | 我们居然更快(但只产了9个chunks vs 53) |
-| d03_tool_call | 2.12s, tool=get_weather | **49.04s, tools=0 finish=stop** | **致命**: 模型不输出<tool_call> XML |
-| d04_reasoning | 17.78s, content[181] reasoning[1011] | 128.74s, content[0] reasoning[1447] | 7x慢; 我们有reasoning但没有content |
-
-## 二、三大根因(按严重程度排序)
-
-### 根因1: GatedDeltaNet每层产NaN → 模型"智力"丧失
-
-docker日志证据:
-```
-WARNING qwen3_5.py:445] NaN in prefill GatedDeltaNet layer 0 (frac=0.9998)
-WARNING qwen3_5.py:445] NaN in prefill GatedDeltaNet layer 1 (frac=0.9997)
-WARNING qwen3_5.py:445] NaN in prefill GatedDeltaNet layer 2 (frac=1.0000)
-WARNING qwen3_5.py:445] NaN in prefill GatedDeltaNet layer 4 (frac=1.0000)
-```
-
-**99.98%-100% NaN率**。`nan_to_num(result, nan=0.0)` 将这些NaN替换为零，等于整个DeltaNet层输出全是零。
-这是一种"活着但脑死亡"的状态——前向传播不报错，但模型失去了DeltaNet层的能力。
-
-**NaN来源追踪**:
-1. `_torch_chunk_gated_delta_rule` 中 `g.cumsum(dim=-1)` → 累积值可能极大
-2. `g.clamp(-20,20)` 后 `g.exp()` → 最大 ~5e8，但这些值进入矩阵乘法后仍可能溢出
-3. `decay_mask = (g_diff).tril().exp()` → 即使单个exp不溢出，大矩阵乘法的累加也可能溢出
-4. `_forward_sub_lower` 中的前向替代: `x[i] = rhs[i] + A[i,:i] @ x[:i]`，如果A中有大值，误差逐行放大
-
-**对手为什么没有这个问题**: 对手可能用的是不同的模型架构(不含DeltaNet)，或者在NVIDIA GPU上float32精度够高不会溢出。
-
-### 根因2: FusedMoE完全fallback → 性能灾难
-
-```
-ERROR _custom_ops.py:58] module 'ixformer.functions' has no attribute 'vllm_moe_topk_softmax'
-WARNING qwen3_5.py:913] FusedMoE native kernel failed, falling back to pure PyTorch experts permanently.
-```
-
-BI-V100的ixformer没有MoE kernel，所有MoE层都用纯PyTorch:
-- 256个expert × top_k=8 → 最多256次F.linear调用(prefill)
-- 每次decode也需要top_k=8次expert forward
-- 对比native kernel的1次fused launch，这是数量级的差距
-
-### 根因3: computility-run.yaml vs 实际参数不一致
-
-yaml写的: `--max-model-len 256000 --max-num-seqs 2 --gpu-memory-utilization 0.95`
-docker日志: `max_seq_len=100000, max_num_seqs=1, gpu_memory_utilization=0.9`
-
-**可能原因**: 部署时还在用旧的配置。需要确认yaml是否真的被用于部署。
-
-## 三、d03_tool_call为什么FAIL
-
-d03日志: `tools=0 finish=stop reasoning[0] (tool_choice=auto) (49.04s)`
-
-**reasoning[0]说明enable_thinking=False确实生效了**。但模型仍然不输出`<tool_call>` XML。
-
-analysis:
-1. enable_thinking=False → 模型不产生`<think>...</think>`块 ✓
-2. 但模型的输出内容不包含`<tool_call><function=get_weather>...` 格式
-3. tool parser `Qwen3CoderToolParser` 在输出中找不到 `<function=` → tools_called=False
-4. 49.04s意味着模型在漫长生成纯文本回答(可能是口头描述天气而不是调用tool)
-
-**核心问题: GatedDeltaNet的NaN导致模型质量太差,不能正确follow tool_call格式**
-
-这不是serving_chat.py的问题。serving_chat.py和protocol.py中的tool_call thinking禁用逻辑是正确的。问题在模型本身。
-
-## 四、对手Sub168分析
-
-对手最终得分60194.6:
-- functional: ~48/52 PASS
-- case_truncation: score=1.0
-- replay_tencent: score=60194 (94/881成功, tps avg 11.86)
-- opencompass: 0.0 (server也崩了)
-
-**对手server在replay后期也崩溃了**(704个connection refused)。
-**对手的replay也只有94/881成功(10.7%)**。
-
-但对手赢在:
-1. functional高通过率 → 基础分
-2. case_truncation通过 → 引擎稳定
-3. replay中94个成功请求 × tps → 得到分数
-
-## 五、修复路径(按投入产出比排序)
-
-### 修复1: NaN问题 — 强制float32精度 + 更激进的clamp
-
-当前: `g.clamp(-20, 20)` 不够。cumsum后再clamp太晚了。
-需要: 在cumsum之前就对g的原始值做clamp。
-
-在 `_torch_chunk_gated_delta_rule`:
-```python
-# 现在: g = g.cumsum(dim=-1).clamp(-20, 20)
-# 改为: g = g.clamp(-5, 5).cumsum(dim=-1).clamp(-15, 15)
-```
-
-在 GatedDeltaNet.forward 的 prefill path:
-```python
-# 现在: _A_safe = self.A_log.float().clamp(-20.0, 20.0)
-# 改为更窄: _A_safe = self.A_log.float().clamp(-10.0, 10.0)
-```
-
-### 修复2: MoE性能 — 尝试真正使用native kernel
-
-docker日志说 `vllm_moe_topk_softmax` 不存在，但 _custom_ops.py 里应该有PyTorch fallback。
-问题是 `_hw_policy.moe_native_align` 和 `_hw_policy.moe_native_invoke` 也是False。
-如果这两个真的不存在，那PyTorch fallback就是唯一选择。
-
-**性能改进**: 在 `_pure_pytorch_experts` 的 prefill path 中:
-- 现在: for-loop over experts, 每个一次F.linear
-- 改为: 按expert batch size排序，大batch的expert合并成一个大F.linear (CCCL histogram pattern已经实现了，但可以更激进)
-
-### 修复3: computility-run.yaml 参数对齐
-
-确保部署时真的用了yaml里的参数。max-num-seqs=2让n=2请求不会崩溃。
-
-### 修复4: d01/d04速度
-
-d01: 95.85s产了1085个token，约11.3 tok/s — 其实tps不太差
-对手d01: 8.49s产了139个token，约16.4 tok/s
-
-**关键差异不是tps，是产了多少token!** 我们1085 vs 对手139。
-我们的模型在thinking里产了大量token。
-d01是basic_nostream，没有tool，所以enable_thinking=True是默认的。
-thinking产了1102个reasoning token + 0个content token。
-
-**问题: d01测试的content[0]意味着没有实际内容输出!**
-对手content[11]说明他输出了内容。
-
-这又回到了GatedDeltaNet NaN → 模型质量差的问题。
-
-## 六、CCCL启示
-
-CCCL在处理数值稳定性方面的核心设计:
-1. `overflow_cast_t<T>` — 在可能溢出的地方用更高精度的中间类型
-2. `cc_dispatch` — 不同硬件不同策略，不硬编码
-3. `policy_selector` — 基于benchmark数据选择参数，不拍脑袋
-
-我们的DeltaNet实现缺少CCCL级别的数值稳定性保证。

SUB509_DIAGNOSIS.md

@@ -1,48 +0,0 @@
-# Sub508/509 完整诊断报告
-
-## 修复提交记录
-
-| Commit | 修复 | 影响 |
-|--------|------|------|
-| e0344b1 | 禁用 tool_call 请求的 thinking | d03 FAIL → 预计 PASS |
-| c241764 | get_scheduler_config try-catch | 防止引擎崩溃 |
-| 994c657 | clamp n>1 to 1 | 防止 t2_n_2 级联崩溃 (19 个测试) |
-
-## Sub508 完整测试结果 (56 tests)
-
-### 实际结果: PASS=21, FAIL=30, SKIP=5
-
-### 级联崩溃 (19 个 FAIL 来自 t2_n_2 引擎崩溃)
-t2_n_2 → HTTP 500 → 引擎死亡 → t3_max_tokens_none/1/64/mid/max/neg1/over,
-t4a/4b, t5, t6, t7, t8, t9, t10, t12_chinese/japanese/emoji 全部 HTTP 500
-
-### 修复后预期: PASS ≈ 40+, FAIL ≈ 10-
-
-### 真正的功能性 FAIL (非级联)
-
-| 测试 | 状态 | 根因 | 可修 |
-|------|------|------|------|
-| d03_tool_call | tools=0 finish=stop | ✅ 已修复 thinking budget | 是 |
-| d05_multimodal | HTTP 400 | multimodal 请求格式 | 需查 |
-| d07_reasoning+content | content[0] | 模型 think 后不产 content | 否(模型) |
-| d10_thinking_disable_ctk | 乱码 content | 模型质量 | 否(模型) |
-| t1a_thinking_true | reasoning[0] | 模型跳过 thinking | 否(模型) |
-| t1c_thinking_default | reasoning[0] | 同上 | 否(模型) |
-| t2_n_2 | HTTP 500 → cascade | ✅ 已修复 clamp n | 是(防崩) |
-
-## 对手 Sub168 对比
-
-| 维度 | 对手 | 我们 |
-|------|------|------|
-| functional PASS | ~50/56 | 21/56 → 修后 ~40/56 |
-| d01 速度 | 8.49s | 95.87s |
-| d04 速度 | 17.78s | 129.19s |
-| replay max_completion_tokens | ✗ 400 rejected (30+次) | ✓ 已支持 (extra=ignore) |
-| replay tool_calls content=None | ✗ 400 rejected | ✓ 已支持 (normalize) |
-| decode TPS | ~16 tok/s | ~11 tok/s |
-
-## 我们 vs 对手的优势
-1. `max_completion_tokens` 支持 — 对手 replay 有 30+ 个 400 错误
-2. `tool_calls` content=None 支持 — 对手 replay preflight 失败
-3. `reasoning_effort` 字段容忍 — 对手被拒
-4. prefix caching 工作 (d06 PASS) — 对手 d06 FAIL

audit_so_usage.sh

@@ -0,0 +1,106 @@
+#!/bin/bash
+set -euo pipefail
+cat << 'PYEOF' | CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" python3 -u -
+"""Audit: which .so functions are actually called in the hot path vs available but unused."""
+import importlib.util, os, sys
+
+SO_DIR = "qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10"
+QWEN = "qwen3_6_scripts/qwen3_5.py"
+PATCH = "qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py"
+XLLM_OPS = "qwen3_6_scripts/ex_engine/python/xllm_ops.py"
+
+# 1. Collect all exported functions from all .so
+print("=" * 70)
+print("  AUDIT: .so function usage")
+print("=" * 70)
+
+so_exports = {}
+for f in sorted(os.listdir(SO_DIR)):
+    if not f.endswith(".so"):
+        continue
+    name = f[:-3]
+    path = os.path.join(SO_DIR, f)
+    try:
+        spec = importlib.util.spec_from_file_location(name, path)
+        m = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(m)
+        fns = [x for x in dir(m) if not x.startswith("_")]
+        so_exports[name] = fns
+    except Exception as e:
+        so_exports[name] = [f"LOAD_ERROR: {e}"]
+
+# 2. Search for usage in qwen3_5.py, patch_vllm_hot_path.py, xllm_ops.py
+code_files = {}
+for label, path in [("qwen3_5.py", QWEN), ("patch_hot_path.py", PATCH), ("xllm_ops.py", XLLM_OPS)]:
+    try:
+        with open(path) as f:
+            code_files[label] = f.read()
+    except:
+        code_files[label] = ""
+
+# Also scan all ex_engine python files
+for f in os.listdir("qwen3_6_scripts/ex_engine/python"):
+    if f.endswith(".py"):
+        path = os.path.join("qwen3_6_scripts/ex_engine/python", f)
+        try:
+            with open(path) as fh:
+                code_files[f"ex_engine/{f}"] = fh.read()
+        except:
+            pass
+
+all_code = "\n".join(code_files.values())
+
+# 3. For each .so and function, check if it's referenced
+print(f"\n{'SO Module':<35} {'Function':<30} {'Used?':<6} {'Where'}")
+print("-" * 110)
+
+total_fns = 0
+used_fns = 0
+unused = []
+
+for so_name in sorted(so_exports.keys()):
+    fns = so_exports[so_name]
+    for fn in fns:
+        if "LOAD_ERROR" in fn:
+            print(f"{so_name:<35} {fn}")
+            continue
+        total_fns += 1
+        
+        # Search patterns: module.fn, .fn(, "fn"
+        found_in = []
+        for label, code in code_files.items():
+            if f".{fn}" in code or f'"{fn}"' in code or f"'{fn}'" in code:
+                found_in.append(label)
+        
+        is_used = len(found_in) > 0
+        if is_used:
+            used_fns += 1
+        else:
+            unused.append((so_name, fn))
+        
+        where = ", ".join(found_in[:3]) if found_in else ""
+        marker = "  ✓" if is_used else "  ✗"
+        print(f"{so_name:<35} {fn:<30} {marker:<6} {where}")
+
+print(f"\n{'=' * 70}")
+print(f"  TOTAL: {used_fns}/{total_fns} functions used")
+print(f"  UNUSED: {total_fns - used_fns} functions")
+print(f"{'=' * 70}")
+
+if unused:
+    print(f"\n  === UNUSED FUNCTIONS ===")
+    for so_name, fn in unused:
+        print(f"    {so_name}.{fn}")
+
+# 4. Check which ixformer_torch_ext functions exist but aren't wrapped
+print(f"\n  === ixformer_torch_ext available but not in any bridge .so ===")
+ix_fns = [
+    "ixformer_linear", "ixformer_linear_ex", "ixformer_linear_allreduce",
+    "linear_i8w8o32", "quantized_linear_awq", "quantized_linear_gptq",
+    "quantized_linear_int8", "quantized_linear_float4", "ixformer_quantized_linear",
+    "silu_and_mul_forward", "rms_norm_forward", "fused_add_rms_norm_forward",
+]
+for fn in ix_fns:
+    in_bridge = fn in all_code
+    print(f"    {fn:<40} {'✓ wrapped' if in_bridge else '✗ NOT wrapped'}")
+PYEOF

bench_gemm.py

@@ -0,0 +1,203 @@
+"""bench_gemm.py — Benchmark all GEMM backends on real device.
+
+Tests with Qwen3.5-27B MoE shapes:
+  - Decode: M=1, K=3584, N=18944*2 (gate_up) / N=3584 (down)
+  - Prefill: M=variable, same K/N
+
+Usage:
+    python3 bench_gemm.py
+"""
+import sys
+import os
+import time
+import torch
+
+# Qwen3.5-27B params (per TP=4 partition)
+H = 3584               # hidden_size
+I = 18944 // 4          # intermediate per partition (4736)
+TWO_I = I * 2           # gate + up
+NUM_EXPERTS = 128
+TOPK = 8
+
+WARMUP = 5
+REPEATS = 20
+
+
+def bench_fn(fn, *args, name=""):
+    """Benchmark a function, return ms per call."""
+    for _ in range(WARMUP):
+        fn(*args)
+    torch.cuda.synchronize()
+
+    t0 = time.perf_counter()
+    for _ in range(REPEATS):
+        fn(*args)
+    torch.cuda.synchronize()
+    elapsed = (time.perf_counter() - t0) / REPEATS * 1000
+    print(f"  {name}: {elapsed:.3f} ms")
+    return elapsed
+
+
+def bench_single_gemm(device):
+    """Benchmark single GEMM: (M,K) × (K,N) for various M."""
+    print("\n=== Single GEMM (M,K)×(K,N) ===")
+    for M in [1, 4, 8, 32]:
+        A = torch.randn(M, H, device=device, dtype=torch.float16)
+        B = torch.randn(H, TWO_I, device=device, dtype=torch.float16)
+
+        bench_fn(torch.mm, A, B, name=f"torch.mm M={M} K={H} N={TWO_I}")
+
+        # Try hgemm
+        try:
+            import hgemm
+            bench_fn(hgemm.hgemm, A, B, name=f"hgemm M={M}")
+        except Exception:
+            pass
+
+        # Try ixformer linear
+        try:
+            import ix_moe_bridge as bridge
+            bench_fn(bridge.linear, A, B.t().contiguous(), name=f"ixformer_linear M={M}")
+        except Exception:
+            pass
+
+
+def bench_group_gemm(device):
+    """Benchmark group GEMM with MoE shapes."""
+    print("\n=== Group GEMM (MoE w13 projection) ===")
+
+    # Simulate decode: 1 token → topk=8 experts, each gets ~1 token
+    total_tokens = TOPK
+    expert_counts = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32)
+    # Distribute tokens to first TOPK experts
+    for i in range(TOPK):
+        expert_counts[i] = 1
+
+    input_t = torch.randn(total_tokens, H, device=device, dtype=torch.float16)
+    w13 = torch.randn(NUM_EXPERTS, TWO_I, H, device=device, dtype=torch.float16) * 0.01
+
+    # PyTorch baseline
+    def torch_group_gemm():
+        offset = 0
+        out = torch.zeros(total_tokens, TWO_I, device=device, dtype=torch.float16)
+        for e in range(NUM_EXPERTS):
+            c = expert_counts[e].item()
+            if c <= 0: continue
+            out[offset:offset+c] = torch.mm(input_t[offset:offset+c], w13[e].t())
+            offset += c
+        return out
+
+    bench_fn(torch_group_gemm, name=f"torch.mm loop (decode, {TOPK} experts)")
+
+    # Try gemm_grouped
+    try:
+        import gemm_grouped
+        bench_fn(gemm_grouped.moe_group_gemm, input_t, w13, expert_counts,
+                 name=f"cutlass_grouped (decode, {TOPK} experts)")
+    except Exception as e:
+        print(f"  cutlass_grouped: {e}")
+
+    # Try ix_moe_bridge
+    try:
+        import ix_moe_bridge as bridge
+        bench_fn(bridge.group_gemm, input_t, w13, expert_counts, TWO_I,
+                 name=f"cuinfer_group_gemm (decode, {TOPK} experts)")
+    except Exception as e:
+        print(f"  cuinfer_group_gemm: {e}")
+
+    # Try hgemm
+    try:
+        import hgemm
+        bench_fn(hgemm.moe_expert_gemm, input_t, w13, expert_counts,
+                 name=f"hgemm_expert (decode, {TOPK} experts)")
+    except Exception as e:
+        print(f"  hgemm_expert: {e}")
+
+    # Prefill shape: 32 tokens
+    print("\n=== Group GEMM (MoE w13, prefill M=32) ===")
+    total_pf = 32 * TOPK  # 256
+    expert_counts_pf = torch.zeros(NUM_EXPERTS, device=device, dtype=torch.int32)
+    for i in range(total_pf):
+        expert_counts_pf[i % NUM_EXPERTS] += 1
+    input_pf = torch.randn(total_pf, H, device=device, dtype=torch.float16)
+
+    def torch_group_gemm_pf():
+        offset = 0
+        out = torch.zeros(total_pf, TWO_I, device=device, dtype=torch.float16)
+        for e in range(NUM_EXPERTS):
+            c = expert_counts_pf[e].item()
+            if c <= 0: continue
+            out[offset:offset+c] = torch.mm(input_pf[offset:offset+c], w13[e].t())
+            offset += c
+        return out
+
+    bench_fn(torch_group_gemm_pf, name=f"torch.mm loop (prefill, 256 tokens)")
+
+    try:
+        import gemm_grouped
+        bench_fn(gemm_grouped.moe_group_gemm, input_pf, w13, expert_counts_pf,
+                 name=f"cutlass_grouped (prefill, 256 tokens)")
+    except Exception as e:
+        print(f"  cutlass_grouped: {e}")
+
+
+def bench_decode_fused(device):
+    """Benchmark full MoE decode pipeline."""
+    print("\n=== Full MoE Decode (1 token, topk=8) ===")
+    hidden = torch.randn(1, H, device=device, dtype=torch.float16)
+    w13_sel = torch.randn(TOPK, TWO_I, H, device=device, dtype=torch.float16) * 0.01
+    w2_sel = torch.randn(TOPK, H, I, device=device, dtype=torch.float16) * 0.01
+    topk_w = torch.softmax(torch.randn(TOPK), dim=0).to(device)
+
+    # PyTorch baseline
+    def torch_decode():
+        results = []
+        for k in range(TOPK):
+            gu = torch.mm(hidden, w13_sel[k].t())
+            act = torch.silu(gu[:, :I]) * gu[:, I:]
+            down = torch.mm(act, w2_sel[k].t())
+            results.append(down * topk_w[k])
+        return sum(results)
+
+    bench_fn(torch_decode, name="torch.mm loop")
+
+    try:
+        import gemm_grouped
+        bench_fn(gemm_grouped.moe_decode_cutlass,
+                 hidden, w13_sel, w2_sel, topk_w,
+                 name="cutlass_batched")
+    except Exception as e:
+        print(f"  cutlass_batched: {e}")
+
+    try:
+        import corex_batched_gemm
+        bench_fn(corex_batched_gemm.moe_decode_fused,
+                 hidden, w13_sel, w2_sel, topk_w,
+                 name="corex_batched")
+    except Exception as e:
+        print(f"  corex_batched: {e}")
+
+
+def main():
+    if not torch.cuda.is_available():
+        print("No CUDA, skipping")
+        sys.exit(0)
+
+    device = torch.device("cuda:0")
+    print(f"Device: {torch.cuda.get_device_name(0)}")
+    print(f"Shapes: H={H}, I={I}, 2I={TWO_I}, experts={NUM_EXPERTS}, topk={TOPK}")
+
+    bench_single_gemm(device)
+    bench_group_gemm(device)
+    bench_decode_fused(device)
+
+    print("\n=== Active backend ===")
+    try:
+        from gemm_dispatch import get_backend
+        print(f"  gemm_dispatch: {get_backend()}")
+    except Exception:
+        print("  gemm_dispatch not loaded")
+
+
+if __name__ == "__main__":
+    main()

bench_linear_patch.sh

@@ -0,0 +1,105 @@
+#!/bin/bash
+set -euo pipefail
+cat << 'PYEOF' | CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" python3 -u -
+import torch, importlib.util, time
+torch.cuda.set_device(0)
+dev = torch.device("cuda:0")
+
+SO = "qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10"
+def load_so(name):
+    spec = importlib.util.spec_from_file_location(name, f"{SO}/{name}.so")
+    m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m); return m
+
+bridge = load_so("ix_moe_bridge")
+act_m = load_so("xllm_activation")
+
+H = 2048
+
+def bench(name, fn, N=1000):
+    for _ in range(100): fn()
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N): fn()
+    torch.cuda.synchronize()
+    us = (time.perf_counter() - t0) / N * 1e6
+    print(f"  {name:50s}: {us:8.1f} us")
+    return us
+
+x = torch.randn(1, H, device=dev, dtype=torch.float16)
+
+# Match upstream gemv_conditions: m <= 1, k % 32 == 0, n % 2 == 0, no bias
+# Test ALL linear ops in qwen3_5.py decode path
+
+print("=== Every linear op in one decode step (TP=4) ===")
+print("--- Attention layer (32 layers) ---")
+
+# QKV: (1,2048) @ (1024,2048)^T → (1,1024)  [heads*head_dim + 2*kv_heads*head_dim]
+w_qkv = torch.randn(1024, H, device=dev, dtype=torch.float16) * 0.01
+bench("qkv F.linear (1,2048)→(1,1024)", lambda: torch.nn.functional.linear(x, w_qkv))
+bench("qkv bridge.linear", lambda: bridge.linear(x, w_qkv, None))
+
+# O_proj: (1,768) @ (2048,768)^T → (1,2048)
+x_o = torch.randn(1, 768, device=dev, dtype=torch.float16)
+w_o = torch.randn(H, 768, device=dev, dtype=torch.float16) * 0.01
+bench("o_proj F.linear (1,768)→(1,2048)", lambda: torch.nn.functional.linear(x_o, w_o))
+bench("o_proj bridge.linear", lambda: bridge.linear(x_o, w_o, None))
+
+print("\n--- GDN layer (4 layers) ---")
+# GDN in_proj: (1,2048) @ (3852,2048)^T → (1,3852)
+w_gdn = torch.randn(3852, H, device=dev, dtype=torch.float16) * 0.01
+bench("gdn_proj F.linear (1,2048)→(1,3852)", lambda: torch.nn.functional.linear(x, w_gdn))
+bench("gdn_proj bridge.linear", lambda: bridge.linear(x, w_gdn, None))
+
+# GDN o_proj: (1,1536) @ (2048,1536)^T → (1,2048)
+x_gdn_o = torch.randn(1, 1536, device=dev, dtype=torch.float16)
+w_gdn_o = torch.randn(H, 1536, device=dev, dtype=torch.float16) * 0.01
+bench("gdn_oproj F.linear (1,1536)→(1,2048)", lambda: torch.nn.functional.linear(x_gdn_o, w_gdn_o))
+bench("gdn_oproj bridge.linear", lambda: bridge.linear(x_gdn_o, w_gdn_o, None))
+
+print("\n--- MoE shared expert (36 layers) ---")
+I_shared = 128
+w_gu = torch.randn(2*I_shared, H, device=dev, dtype=torch.float16) * 0.01
+w_down = torch.randn(H, I_shared, device=dev, dtype=torch.float16) * 0.01
+bench("shared gate_up F.linear (1,2048)→(1,256)", lambda: torch.nn.functional.linear(x, w_gu))
+bench("shared gate_up bridge.linear", lambda: bridge.linear(x, w_gu, None))
+x_down = torch.randn(1, I_shared, device=dev, dtype=torch.float16)
+bench("shared down F.linear (1,128)→(1,2048)", lambda: torch.nn.functional.linear(x_down, w_down))
+bench("shared down bridge.linear", lambda: bridge.linear(x_down, w_down, None))
+
+print("\n--- Router (36 layers) ---")
+w_router = torch.randn(257, H, device=dev, dtype=torch.float16) * 0.01
+bench("router F.linear (1,2048)→(1,257)", lambda: torch.nn.functional.linear(x, w_router))
+bench("router bridge.linear", lambda: bridge.linear(x, w_router, None))
+
+print("\n--- LM head (1x) ---")
+w_lm = torch.randn(37984, H, device=dev, dtype=torch.float16) * 0.01
+bench("lm_head F.linear (1,2048)→(1,37984)", lambda: torch.nn.functional.linear(x, w_lm))
+bench("lm_head bridge.linear", lambda: bridge.linear(x, w_lm, None))
+
+# === Total impact ===
+print("\n=== Projected total decode step savings ===")
+shapes = [
+    ("attn_qkv",   32, (1024, H)),
+    ("attn_o",     32, (H, 768)),
+    ("gdn_proj",    4, (3852, H)),
+    ("gdn_o",       4, (H, 1536)),
+    ("shared_gu",  36, (2*I_shared, H)),
+    ("shared_down",36, (H, I_shared)),
+    ("router",     36, (257, H)),
+    ("lm_head",     1, (37984, H)),
+]
+total_torch = 0
+total_bridge = 0
+for name, count, (N, K) in shapes:
+    w = torch.randn(N, K, device=dev, dtype=torch.float16) * 0.01
+    xi = torch.randn(1, K, device=dev, dtype=torch.float16)
+    t_torch = bench(f"  {name} F.linear", lambda xi=xi, w=w: torch.nn.functional.linear(xi, w), N=500)
+    t_bridge = bench(f"  {name} bridge", lambda xi=xi, w=w: bridge.linear(xi, w, None), N=500)
+    total_torch += t_torch * count
+    total_bridge += t_bridge * count
+    speedup = t_torch / t_bridge if t_bridge > 0 else 0
+    print(f"    → x{count}: {t_torch*count:.0f} → {t_bridge*count:.0f} us ({speedup:.1f}x)")
+
+print(f"\n  TOTAL linear ops: {total_torch:.0f} → {total_bridge:.0f} us")
+print(f"  Savings: {total_torch - total_bridge:.0f} us = {(total_torch-total_bridge)/1000:.1f} ms")
+PYEOF

bench_shared.sh

@@ -0,0 +1,90 @@
+#!/bin/bash
+set -euo pipefail
+cat << 'PYEOF' | CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" python3 -u -
+import torch, importlib.util, time
+torch.cuda.set_device(0)
+dev = torch.device("cuda:0")
+
+SO = "qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10"
+def load_so(name):
+    spec = importlib.util.spec_from_file_location(name, f"{SO}/{name}.so")
+    m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m); return m
+
+bridge = load_so("ix_moe_bridge")
+act_m = load_so("xllm_activation")
+
+H = 2048
+I_shared = 128
+
+x = torch.randn(1, H, device=dev, dtype=torch.float16)
+w_gu = torch.randn(2*I_shared, H, device=dev, dtype=torch.float16) * 0.01
+w_down = torch.randn(H, I_shared, device=dev, dtype=torch.float16) * 0.01
+
+def bench(name, fn, N=1000):
+    for _ in range(100): fn()
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(N): fn()
+    torch.cuda.synchronize()
+    us = (time.perf_counter() - t0) / N * 1e6
+    print(f"  {name:45s}: {us:8.1f} us")
+    return us
+
+print("=== ix_moe_bridge.linear probe ===")
+linear_ok = False
+for desc, args in [("(x,w)", (x, w_gu)),
+                    ("(x,w,None)", (x, w_gu, None)),
+                    ("(x,w,bias0)", (x, w_gu, torch.zeros(2*I_shared,device=dev,dtype=torch.float16)))]:
+    try:
+        out = bridge.linear(*args); torch.cuda.synchronize()
+        print(f"  linear{desc}: OK shape={out.shape}")
+        linear_ok = True; break
+    except Exception as e:
+        print(f"  linear{desc}: {str(e)[:80]}")
+
+print("\n=== Shared expert benchmarks ===")
+act_buf = torch.empty(1, I_shared, device=dev, dtype=torch.float16)
+
+def shared_torch():
+    gu = torch.nn.functional.linear(x, w_gu)
+    g, u = gu.chunk(2, dim=-1)
+    act = torch.sigmoid(g) * g * u
+    return torch.nn.functional.linear(act, w_down)
+bench("A: torch linear + torch silu", shared_torch)
+
+def shared_xllm_silu():
+    gu = torch.nn.functional.linear(x, w_gu)
+    act_m.silu_and_mul(act_buf, gu)
+    return torch.nn.functional.linear(act_buf, w_down)
+bench("B: torch linear + xllm silu", shared_xllm_silu)
+
+if linear_ok:
+    try:
+        _ = bridge.linear(x, w_gu)
+        def shared_bridge():
+            gu = bridge.linear(x, w_gu)
+            act_m.silu_and_mul(act_buf, gu)
+            return bridge.linear(act_buf, w_down)
+        bench("C: bridge linear + xllm silu", shared_bridge)
+    except:
+        try:
+            b_gu = torch.zeros(2*I_shared,device=dev,dtype=torch.float16)
+            b_dn = torch.zeros(H,device=dev,dtype=torch.float16)
+            def shared_bridge_b():
+                gu = bridge.linear(x, w_gu, b_gu)
+                act_m.silu_and_mul(act_buf, gu)
+                return bridge.linear(act_buf, w_down, b_dn)
+            bench("C: bridge linear(bias0) + xllm silu", shared_bridge_b)
+        except Exception as e:
+            print(f"  C failed: {e}")
+
+print("\n=== Step breakdown ===")
+bench("gate_up F.linear (1,2048)@(256,2048)^T", lambda: torch.nn.functional.linear(x, w_gu))
+gu_t = torch.nn.functional.linear(x, w_gu)
+bench("silu_and_mul", lambda: act_m.silu_and_mul(act_buf, gu_t))
+bench("down F.linear (1,128)@(2048,128)^T", lambda: torch.nn.functional.linear(act_buf, w_down))
+
+print("\n=== matmul vs F.linear ===")
+bench("torch.mm(x, w_gu.T)", lambda: torch.mm(x, w_gu.t()))
+bench("F.linear(x, w_gu)", lambda: torch.nn.functional.linear(x, w_gu))
+PYEOF

build_moe_bridge.sh

@@ -0,0 +1,179 @@
+#!/usr/bin/env bash
+# build_moe_bridge.sh — Compile MoE ops + bridge into ix_moe_bridge.so
+#
+# Links against:
+#   libcuinfer.so   (cuinferCustomGemm, cuinferTopK — confirmed in symbol dump)
+#   libixformer.so  (silu_and_mul, rms_norm, flash_attn, etc — confirmed)
+#
+# Real device compiler: corex clang/16, NOT nvcc
+# Reference: ex_engine/build_ix_bridge.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+VLLM_ROOT="${1:-}"
+
+echo "[moe_bridge] Building ix_moe_bridge.so"
+echo "[moe_bridge] Script dir: ${SCRIPT_DIR}"
+
+# --- Locate sources ---
+# Support both layouts:
+#   1. SCRIPT_DIR=/workspace/ex_engine → csrc/ is direct child
+#   2. SCRIPT_DIR=/workspace/qwen3_6_scripts/ex_engine_src → csrc/ is direct child
+MOE_CU=""
+BRIDGE_CPP=""
+for base in "${SCRIPT_DIR}" "${SCRIPT_DIR}/ex_engine"; do
+    [[ -f "${base}/csrc/moe_ops_impl.cu" ]] && MOE_CU="${base}/csrc/moe_ops_impl.cu"
+    [[ -f "${base}/csrc/ix_full_bridge_v2.cpp" ]] && BRIDGE_CPP="${base}/csrc/ix_full_bridge_v2.cpp"
+done
+
+if [[ -z "$MOE_CU" ]]; then
+    echo "[moe_bridge] ERROR: moe_ops_impl.cu not found under ${SCRIPT_DIR}" >&2
+    exit 1
+fi
+if [[ -z "$BRIDGE_CPP" ]]; then
+    echo "[moe_bridge] ERROR: ix_full_bridge_v2.cpp not found under ${SCRIPT_DIR}" >&2
+    exit 1
+fi
+echo "[moe_bridge] MOE_CU: ${MOE_CU}"
+echo "[moe_bridge] BRIDGE_CPP: ${BRIDGE_CPP}"
+
+# --- Locate libraries ---
+COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
+
+# Find libcuinfer.so
+CUINFER_SO=""
+for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib" "/usr/lib64" "/usr/lib"; do
+    if [[ -f "${d}/libcuinfer.so" ]]; then
+        CUINFER_SO="${d}/libcuinfer.so"
+        break
+    fi
+done
+
+# Find libixformer.so and ixformer Python package
+IX_LIB_DIR=""
+IX_SO_FILES=()
+for d in \
+    "${COREX_ROOT}/lib/python3/dist-packages/ixformer" \
+    "${COREX_ROOT}/lib64/python3/dist-packages/ixformer" \
+    "$(python3 -c 'import ixformer, os; print(os.path.dirname(ixformer.__file__))' 2>/dev/null || echo '')"; do
+    if [[ -d "$d" ]]; then
+        IX_LIB_DIR="$d"
+        while IFS= read -r so; do
+            IX_SO_FILES+=("$so")
+        done < <(find "$d" -name "*.so" -type f 2>/dev/null)
+        break
+    fi
+done
+
+echo "[moe_bridge] COREX_ROOT: ${COREX_ROOT}"
+echo "[moe_bridge] cuinfer: ${CUINFER_SO:-NOT FOUND}"
+echo "[moe_bridge] ixformer dir: ${IX_LIB_DIR:-NOT FOUND}"
+echo "[moe_bridge] ixformer .so count: ${#IX_SO_FILES[@]}"
+
+# --- Build via torch.utils.cpp_extension ---
+mkdir -p "${SCRIPT_DIR}/prebuilt"
+
+export SCRIPT_DIR VLLM_ROOT
+python3 << 'PYEOF'
+import os, sys, glob, shutil
+
+script_dir = os.environ.get("SCRIPT_DIR", ".")
+vllm_root = os.environ.get("VLLM_ROOT", "")
+
+# Find source files — try direct csrc/ first, then ex_engine/csrc/
+moe_cu = ""
+bridge_cpp = ""
+for base in [script_dir, os.path.join(script_dir, "ex_engine")]:
+    candidate_cu = os.path.join(base, "csrc", "moe_ops_impl.cu")
+    candidate_cpp = os.path.join(base, "csrc", "ix_full_bridge_v2.cpp")
+    if os.path.isfile(candidate_cu):
+        moe_cu = candidate_cu
+    if os.path.isfile(candidate_cpp):
+        bridge_cpp = candidate_cpp
+if not moe_cu or not bridge_cpp:
+    print(f"[moe_bridge] ERROR: sources not found under {script_dir}")
+    sys.exit(1)
+print(f"[moe_bridge] MOE_CU: {moe_cu}")
+print(f"[moe_bridge] BRIDGE_CPP: {bridge_cpp}")
+
+# Collect linker flags
+extra_ldflags = []
+rpath_dirs = set()
+
+corex_root = os.environ.get("COREX_ROOT", "/usr/local/corex")
+for search_dir in [
+    os.path.join(corex_root, "lib64"),
+    os.path.join(corex_root, "lib"),
+]:
+    if os.path.isdir(search_dir):
+        rpath_dirs.add(search_dir)
+        for so in glob.glob(os.path.join(search_dir, "libcuinfer*.so*")):
+            extra_ldflags.append(so)
+
+# ixformer .so files
+try:
+    import ixformer
+    ix_dir = os.path.dirname(ixformer.__file__)
+    rpath_dirs.add(ix_dir)
+    for so in glob.glob(os.path.join(ix_dir, "*.so")):
+        extra_ldflags.append(so)
+    for so in glob.glob(os.path.join(ix_dir, "lib*.so")):
+        if so not in extra_ldflags:
+            extra_ldflags.append(so)
+except ImportError:
+    # Search common paths
+    for d in [
+        os.path.join(corex_root, "lib", "python3", "dist-packages", "ixformer"),
+        os.path.join(corex_root, "lib64", "python3", "dist-packages", "ixformer"),
+    ]:
+        if os.path.isdir(d):
+            rpath_dirs.add(d)
+            for so in glob.glob(os.path.join(d, "*.so")):
+                extra_ldflags.append(so)
+
+for d in rpath_dirs:
+    extra_ldflags.append(f"-Wl,-rpath,{d}")
+
+print(f"[moe_bridge] Linking against {len(extra_ldflags)} items")
+for f in extra_ldflags[:10]:
+    print(f"  {f}")
+
+try:
+    from torch.utils.cpp_extension import load
+
+    mod = load(
+        name="ix_moe_bridge",
+        sources=[moe_cu, bridge_cpp],
+        extra_include_paths=[os.path.join(script_dir, "csrc")],
+        extra_cflags=["-O2", "-std=c++17"],
+        extra_cuda_cflags=["-O2", ],
+        extra_ldflags=extra_ldflags,
+        verbose=True,
+    )
+    print("[moe_bridge] ✓ Compilation successful")
+
+    # Find and copy the built .so
+    import importlib
+    spec = importlib.util.find_spec("ix_moe_bridge")
+    if spec and spec.origin:
+        dst = os.path.join(script_dir, "prebuilt", "ix_moe_bridge.so")
+        shutil.copy2(spec.origin, dst)
+        print(f"[moe_bridge] ✓ Saved to {dst}")
+
+        if vllm_root:
+            vllm_dst = os.path.join(vllm_root, "ex_engine", "ix_moe_bridge.so")
+            os.makedirs(os.path.dirname(vllm_dst), exist_ok=True)
+            shutil.copy2(spec.origin, vllm_dst)
+            print(f"[moe_bridge] ✓ Deployed to {vllm_dst}")
+    else:
+        print("[moe_bridge] ⚠ Could not locate compiled .so via importlib")
+
+except Exception as e:
+    print(f"[moe_bridge] ERROR: {e}", file=sys.stderr)
+    import traceback; traceback.print_exc()
+    sys.exit(1)
+PYEOF
+
+echo "[moe_bridge] Done"

cat_files/CMakeLists_batched_gemm.txt

@@ -0,0 +1,27 @@
+# Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without modification, are permitted
+# provided that the following conditions are met:
+#     * Redistributions of source code must retain the above copyright notice, this list of
+#       conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above copyright notice, this list of
+#       conditions and the following disclaimer in the documentation and/or other materials
+#       provided with the distribution.
+#     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+#       to endorse or promote products derived from this software without specific prior written
+#       permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+# STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+cutlass_example_add_executable(
+  05_batched_gemm
+  batched_gemm.cu
+)
+  

cat_files/CMakeLists_tensorop_gemm.txt

@@ -0,0 +1,27 @@
+# Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without modification, are permitted
+# provided that the following conditions are met:
+#     * Redistributions of source code must retain the above copyright notice, this list of
+#       conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above copyright notice, this list of
+#       conditions and the following disclaimer in the documentation and/or other materials
+#       provided with the distribution.
+#     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+#       to endorse or promote products derived from this software without specific prior written
+#       permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+# STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+# cutlass_example_add_executable(
+#   08_turing_tensorop_gemm
+#   turing_tensorop_gemm.cu
+#   )
+

cat_files/arch.h

@@ -0,0 +1,84 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ *modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *notice, this list of conditions and the following disclaimer in the
+ *documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its
+ *contributors may be used to endorse or promote products derived from this
+ *software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ *AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ *IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ *DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY DIRECT,
+ *INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ *DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ *OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TOR (INCLUDING
+ *NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+ *EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Defines tags for architecture-specific configurations.
+*/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace arch {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+struct Sm50 {
+    static int const kMinComputeCapability = 50;
+};
+struct Sm60 {
+    static int const kMinComputeCapability = 60;
+};
+struct Sm61 {
+    static int const kMinComputeCapability = 61;
+};
+
+
+/// BIGISLAND Arch
+struct Cu10 {
+    static int const kMinComputeCapability = 10;
+};
+
+struct Sm62 {
+    static int const kMinComputeCapability = 62;
+};
+
+/// Triggers a breakpoint on the device
+CUTLASS_DEVICE
+void device_breakpoint() {
+#if defined(__CUDA_ARCH__)
+  asm volatile ("  brkpt;\n");
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// Switches to control performance improvement only supported on Iluvatar platform
+
+/// Compiler of Iluvatar-CoreX implicitly convert boolean type that is stored at VRF to a 64-bit 
+/// width integer type on SRF
+#define IMPLICIT_VRF_BOOLEAN_TO_SRF_INTEGER 1
+
+/// Enable block load or store
+#define BLOCK_LOAD_STORE 1
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace arch
+}  // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////

cat_files/arch_arch.h

@@ -0,0 +1,84 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ *modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *notice, this list of conditions and the following disclaimer in the
+ *documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its
+ *contributors may be used to endorse or promote products derived from this
+ *software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ *AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ *IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ *DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY DIRECT,
+ *INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ *DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ *OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TOR (INCLUDING
+ *NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+ *EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Defines tags for architecture-specific configurations.
+*/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace arch {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+struct Sm50 {
+    static int const kMinComputeCapability = 50;
+};
+struct Sm60 {
+    static int const kMinComputeCapability = 60;
+};
+struct Sm61 {
+    static int const kMinComputeCapability = 61;
+};
+
+
+/// BIGISLAND Arch
+struct Cu10 {
+    static int const kMinComputeCapability = 10;
+};
+
+struct Sm62 {
+    static int const kMinComputeCapability = 62;
+};
+
+/// Triggers a breakpoint on the device
+CUTLASS_DEVICE
+void device_breakpoint() {
+#if defined(__CUDA_ARCH__)
+  asm volatile ("  brkpt;\n");
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// Switches to control performance improvement only supported on Iluvatar platform
+
+/// Compiler of Iluvatar-CoreX implicitly convert boolean type that is stored at VRF to a 64-bit 
+/// width integer type on SRF
+#define IMPLICIT_VRF_BOOLEAN_TO_SRF_INTEGER 1
+
+/// Enable block load or store
+#define BLOCK_LOAD_STORE 1
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace arch
+}  // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////

cat_files/basic_gemm.cu

@@ -0,0 +1,492 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*
+  This example demonstrates how to call a CUTLASS GEMM kernel and provides a naive reference
+  matrix multiply kernel to verify its correctness.
+
+  The CUTLASS Gemm template is instantiated in the function CutlassSgemmNN. This is kernel computes
+  the general matrix product (GEMM) using single-precision floating-point arithmetic and assumes
+  all matrices have column-major layout.
+
+  The threadblock tile size is chosen as 128x128x8 which offers good performance for large matrices.
+  See the CUTLASS Parallel for All blog post for more exposition on the tunable parameters available
+  in CUTLASS.
+
+  https://devblogs.nvidia.com/cutlass-linear-algebra-cuda/
+
+  Aside from defining and launching the SGEMM kernel, this example does not use any other components
+  or utilities within CUTLASS. Such utilities are demonstrated elsewhere in other examples and are
+  prevalent in the CUTLASS unit tests.
+
+  This example has delibrately been kept similar to the basic_gemm example from cutass-1.3 to 
+  highlight the minimum amount of differences needed to transition to cutlass-2.0.
+
+  Cutlass-1.3 sgemm: https://github.com/NVIDIA/cutlass/blob/master/examples/00_basic_gemm/basic_gemm.cu
+*/
+
+// Standard Library includes
+#include <iostream>
+#include <sstream>
+#include <vector>
+
+// Helper methods to check for errors
+#include "helper.h"
+
+//
+// CUTLASS includes needed for single-precision GEMM kernel
+//
+
+// Defines cutlass::gemm::device::Gemm, the generic Gemm computation template class.
+#include "cutlass/gemm/device/gemm.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// This function defines a CUTLASS GEMM kernel instantiation, constructs its parameters object,
+// and launches it on the CUDA device.
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Define a CUTLASS GEMM template and launch a GEMM kernel.
+cudaError_t CutlassSgemmNN(
+  int M,
+  int N,
+  int K,
+  float alpha,
+  float const *A,
+  int lda,
+  float const *B,
+  int ldb,
+  float beta,
+  float *C,
+  int ldc) {
+
+  // Define type definition for single-precision CUTLASS GEMM with column-major
+  // input matrices and 128x128x8 threadblock tile size (chosen by default).
+  //
+  // To keep the interface manageable, several helpers are defined for plausible compositions
+  // including the following example for single-precision GEMM. Typical values are used as
+  // default template arguments. See `cutlass/gemm/device/default_gemm_configuration.h` for more details.
+  //
+  // To view the full gemm device API interface, see `cutlass/gemm/device/gemm.h`
+
+  using ColumnMajor = cutlass::layout::ColumnMajor;
+
+  using CutlassGemm = cutlass::gemm::device::Gemm<float,        // Data-type of A matrix
+                                                  ColumnMajor,  // Layout of A matrix
+                                                  float,        // Data-type of B matrix
+                                                  ColumnMajor,  // Layout of B matrix
+                                                  float,        // Data-type of C matrix
+                                                  ColumnMajor>; // Layout of C matrix
+
+  // Define a CUTLASS GEMM type
+  CutlassGemm gemm_operator;
+
+  // Construct the CUTLASS GEMM arguments object.
+  //
+  // One of CUTLASS's design patterns is to define gemm argument objects that are constructible
+  // in host code and passed to kernels by value. These may include pointers, strides, scalars,
+  // and other arguments needed by Gemm and its components.
+  //
+  // The benefits of this pattern are (1.) a structured, composable strategy for passing host-constructible
+  // arguments to kernels and (2.) minimized initialization overhead on kernel entry.
+  //
+  CutlassGemm::Arguments args({M , N, K},  // Gemm Problem dimensions
+                              {A, lda},    // Tensor-ref for source matrix A
+                              {B, ldb},    // Tensor-ref for source matrix B
+                              {C, ldc},    // Tensor-ref for source matrix C
+                              {C, ldc},    // Tensor-ref for destination matrix D (may be different memory than source C matrix)
+                              {alpha, beta}); // Scalars used in the Epilogue
+
+  //
+  // Launch the CUTLASS GEMM kernel.
+  //
+  
+  cutlass::Status status = gemm_operator(args);
+
+  //
+  // Return a cudaError_t if the CUTLASS GEMM operator returned an error code.
+  //
+
+  if (status != cutlass::Status::kSuccess) {
+    return cudaErrorUnknown;
+  }
+
+  // Return success, if no errors were encountered.
+  return cudaSuccess;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// The source code after this point in the file is generic CUDA using the CUDA Runtime API
+// and simple CUDA kernels to initialize matrices and compute the general matrix product.
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Kernel to initialize a matrix with small integers.
+__global__ void InitializeMatrix_kernel(
+  float *matrix,
+  int ldm,
+  int rows,
+  int columns,
+  int seed = 0) {
+
+  int i = threadIdx.x + blockIdx.x * blockDim.x;
+  int j = threadIdx.y + blockIdx.y * blockDim.y;
+
+  if (i < rows && j < columns) {
+    int offset = i + j * ldm;
+
+    // Generate arbitrary elements.
+    int const k = 16807;
+    int const m = 16;
+    float value = float(((offset + seed) * k % m) - m / 2);
+
+    matrix[offset] = value;
+  }
+}
+
+/// Simple function to initialize a matrix to arbitrary small integers.
+cudaError_t InitializeMatrix(float *matrix, int ldm, int rows, int columns, int seed = 0) {
+
+  dim3 block(16, 16);
+  dim3 grid(
+    (rows + block.x - 1) / block.x,
+    (columns + block.y - 1) / block.y
+  );
+
+  InitializeMatrix_kernel<<< grid, block >>>(matrix, ldm, rows, columns, seed);
+
+  return cudaGetLastError();
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Allocates device memory for a matrix then fills with arbitrary small integers.
+cudaError_t AllocateMatrix(float **matrix, int ldm, int rows, int columns, int seed = 0) {
+  cudaError_t result;
+
+  size_t sizeof_matrix = sizeof(float) * ldm * columns;
+
+  // Allocate device memory.
+  result = cudaMalloc(reinterpret_cast<void **>(matrix), sizeof_matrix);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to allocate matrix: "
+      << cudaGetErrorString(result) << std::endl;
+    return result;
+  }
+
+  // Clear the allocation.
+  result = cudaMemset(*matrix, 0, sizeof_matrix);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to clear matrix device memory: "
+      << cudaGetErrorString(result) << std::endl;
+    return result;
+  }
+
+  // Initialize matrix elements to arbitrary small integers.
+  result = InitializeMatrix(*matrix, ldm, rows, columns, seed);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to initialize matrix: "
+      << cudaGetErrorString(result) << std::endl;
+    return result;
+  }
+
+  return result;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Naive reference GEMM computation.
+__global__ void ReferenceGemm_kernel(
+  int M,
+  int N,
+  int K,
+  float alpha,
+  float const *A,
+  int lda,
+  float const *B,
+  int ldb,
+  float beta,
+  float *C,
+  int ldc) {
+
+  int i = threadIdx.x + blockIdx.x * blockDim.x;
+  int j = threadIdx.y + blockIdx.y * blockDim.y;
+
+  if (i < M && j < N) {
+    float accumulator = 0;
+
+    for (int k = 0; k < K; ++k) {
+      accumulator += A[i + k * lda] * B[k + j * ldb];
+    }
+
+    C[i + j * ldc] = alpha * accumulator + beta * C[i + j * ldc];
+  }
+}
+
+/// Reference GEMM computation.
+cudaError_t ReferenceGemm(
+  int M,
+  int N,
+  int K,
+  float alpha,
+  float const *A,
+  int lda,
+  float const *B,
+  int ldb,
+  float beta,
+  float *C,
+  int ldc) {
+
+  dim3 block(16, 16);
+  dim3 grid(
+    (M + block.x - 1) / block.x,
+    (N + block.y - 1) / block.y
+  );
+
+  ReferenceGemm_kernel<<< grid, block >>>(M, N, K, alpha, A, lda, B, ldb, beta, C, ldc);
+
+  return cudaGetLastError();
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Allocate several matrices in GPU device memory and call a single-precision
+/// CUTLASS GEMM kernel.
+cudaError_t TestCutlassGemm(int M, int N, int K, float alpha, float beta) {
+  cudaError_t result;
+
+  //
+  // Define several matrices to be used as operands to GEMM kernels.
+  //
+
+  // Compute leading dimensions for each matrix.
+  int lda = M;
+  int ldb = K;
+  int ldc = M;
+
+  // Compute size in bytes of the C matrix.
+  size_t sizeof_C = sizeof(float) * ldc * N;
+
+  // Define pointers to matrices in GPU device memory.
+  float *A;
+  float *B;
+  float *C_cutlass;
+  float *C_reference;
+
+  //
+  // Allocate matrices in GPU device memory with arbitrary seeds.
+  //
+
+  result = AllocateMatrix(&A, lda, M, K, 0);
+
+  if (result !=  cudaSuccess) {
+    return result;
+  }
+
+  result = AllocateMatrix(&B, ldb, K, N, 17);
+
+  if (result !=  cudaSuccess) {
+    cudaFree(A);
+    return result;
+  }
+
+  result = AllocateMatrix(&C_cutlass, ldc, M, N, 101);
+
+  if (result != cudaSuccess) {
+    cudaFree(A);
+    cudaFree(B);
+    return result;
+  }
+
+  result = AllocateMatrix(&C_reference, ldc, M, N, 101);
+
+  if (result != cudaSuccess) {
+    cudaFree(A);
+    cudaFree(B);
+    cudaFree(C_cutlass);
+    return result;
+  }
+
+  result = cudaMemcpy(C_reference, C_cutlass, sizeof_C, cudaMemcpyDeviceToDevice);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to copy C_cutlass matrix to C_reference: "
+      << cudaGetErrorString(result) << std::endl;
+
+    cudaFree(C_reference);
+    cudaFree(C_cutlass);
+    cudaFree(B);
+    cudaFree(A);
+
+    return result;
+  }
+
+  //
+  // Launch CUTLASS GEMM.
+  //
+
+  result = CutlassSgemmNN(M, N, K, alpha, A, lda, B, ldb, beta, C_cutlass, ldc);
+
+  if (result != cudaSuccess) {
+    std::cerr << "CUTLASS GEMM kernel failed: "
+      << cudaGetErrorString(result) << std::endl;
+
+    cudaFree(C_reference);
+    cudaFree(C_cutlass);
+    cudaFree(B);
+    cudaFree(A);
+
+    return result;
+  }
+
+  //
+  // Verify.
+  //
+
+  // Launch reference GEMM
+  result = ReferenceGemm(M, N, K, alpha, A, lda, B, ldb, beta, C_reference, ldc);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Reference GEMM kernel failed: "
+      << cudaGetErrorString(result) << std::endl;
+
+    cudaFree(C_reference);
+    cudaFree(C_cutlass);
+    cudaFree(B);
+    cudaFree(A);
+
+    return result;
+  }
+
+  // Copy to host and verify equivalence.
+  std::vector<float> host_cutlass(ldc * N, 0);
+  std::vector<float> host_reference(ldc * N, 0);
+
+  result = cudaMemcpy(host_cutlass.data(), C_cutlass, sizeof_C, cudaMemcpyDeviceToHost);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to copy CUTLASS GEMM results: "
+      << cudaGetErrorString(result) << std::endl;
+
+    cudaFree(C_reference);
+    cudaFree(C_cutlass);
+    cudaFree(B);
+    cudaFree(A);
+
+    return result;
+  }
+
+  result = cudaMemcpy(host_reference.data(), C_reference, sizeof_C, cudaMemcpyDeviceToHost);
+
+  if (result != cudaSuccess) {
+    std::cerr << "Failed to copy Reference GEMM results: "
+      << cudaGetErrorString(result) << std::endl;
+
+    cudaFree(C_reference);
+    cudaFree(C_cutlass);
+    cudaFree(B);
+    cudaFree(A);
+
+    return result;
+  }
+
+  //
+  // Free device memory allocations.
+  //
+
+  cudaFree(C_reference);
+  cudaFree(C_cutlass);
+  cudaFree(B);
+  cudaFree(A);
+
+  //
+  // Test for bit equivalence of results.
+  //
+
+  if (host_cutlass != host_reference) {
+    std::cerr << "CUTLASS results incorrect." << std::endl;
+
+    return cudaErrorUnknown;
+  }
+
+  return cudaSuccess;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Entry point to basic_gemm example.
+//
+// usage:
+//
+//   00_basic_gemm <M> <N> <K> <alpha> <beta>
+//
+int main(int argc, const char *arg[]) {
+
+  //
+  // Parse the command line to obtain GEMM dimensions and scalar values.
+  //
+
+  // GEMM problem dimensions.
+  int problem[3] = { 128, 128, 128 };
+
+  for (int i = 1; i < argc && i < 4; ++i) {
+    std::stringstream ss(arg[i]);
+    ss >> problem[i - 1];
+  }
+
+  // Scalars used for linear scaling the result of the matrix product.
+  float scalars[2] = { 1, 0 };
+
+  for (int i = 4; i < argc && i < 6; ++i) {
+    std::stringstream ss(arg[i]);
+    ss >> scalars[i - 4];
+  }
+
+  //
+  // Run the CUTLASS GEMM test.
+  //
+
+  cudaError_t result = TestCutlassGemm(
+    problem[0],     // GEMM M dimension
+    problem[1],     // GEMM N dimension
+    problem[2],     // GEMM K dimension
+    scalars[0],     // alpha
+    scalars[1]      // beta
+  );
+
+  if (result == cudaSuccess) {
+    std::cout << "Passed." << std::endl;
+  }
+
+  // Exit.
+  return result == cudaSuccess ? 0 : -1;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////

cat_files/batched_gemm.cu

@@ -0,0 +1,345 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#include <iostream>
+#include <vector>
+
+#include "cutlass/cutlass.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/gemm/device/gemm_batched.h"
+
+#pragma warning( disable : 4503)
+
+/*
+This example demonstrates how to use cutlass to compute a batched strided gemm.
+In this example, both A and B matrix are non-transpose and column major matrix
+batched_C = batched_A x batched_B
+As an example, matrix C can be seen as
+-----------------------------------------------------------
+(0,0,0) | (0,0,1) | (0,0,2) | (1,0,0) | (1,0,1) | (1,0,2) |
+-----------------------------------------------------------
+(0,1,0) | (0,1,1) | (0,1,2) | (1,1,0) | (1,1,1) | (1,1,2) |
+-----------------------------------------------------------
+(0,2,0) | (0,2,1) | (0,2,2) | (1,2,0) | (1,2,1) | (1,2,2) |
+-----------------------------------------------------------
+(0,3,0) | (0,3,1) | (0,3,2) | (1,3,0) | (1,3,1) | (1,3,2) |
+-----------------------------------------------------------
+(0,4,0) | (0,4,1) | (0,4,2) | (1,4,0) | (1,4,1) | (1,4,2) |
+-----------------------------------------------------------
+(0,5,0) | (0,5,1) | (0,5,2) | (1,5,0) | (1,5,1) | (1,5,2) |
+-----------------------------------------------------------
+           batch 0          |           batch 1
+where we denote each element with (batch_idx, row_idx, column_idx)
+In this example, batch size is 2, M is 6 and N is 3
+The stride (batch_stride_C) between the first element of two batches is ldc * n
+
+matrix A can be seen as
+---------------------------------------
+(0,0,0) | (0,0,1) | (1,0,0) | (1,0,1) |
+---------------------------------------
+(0,1,0) | (0,1,1) | (1,1,0) | (1,1,1) |
+---------------------------------------
+(0,2,0) | (0,2,1) | (1,2,0) | (1,2,1) |
+---------------------------------------
+(0,3,0) | (0,3,1) | (1,3,0) | (1,3,1) |
+---------------------------------------
+(0,4,0) | (0,4,1) | (1,4,0) | (1,4,1) |
+---------------------------------------
+(0,5,0) | (0,5,1) | (1,5,0) | (1,5,1) |
+---------------------------------------
+     batch 0      |      batch 1
+, where batch size is 2, M is 6 and K is 2
+The stride (batch_stride_B) between the first element of two batches is lda * k
+
+matrix B can be seen as
+-----------------------------
+(0,0,0) | (0,0,1) | (0,0,2) |
+----------------------------- batch 0
+(0,1,0) | (0,1,1) | (0,1,2) |
+-------------------------------------
+(1,0,0) | (1,0,1) | (1,0,2) |
+----------------------------- batch 1
+(1,1,0) | (1,1,1) | (1,1,2) |
+-----------------------------
+, where the batch size is 2, N is 3 and K is 2
+The stride (batch_stride_C) between the first element of two batches is k
+
+
+*/
+
+cudaError_t cutlass_strided_batched_sgemm(
+  int m, 
+  int n,
+  int k,
+  float alpha,
+  float const *A,
+  int lda,
+  long long int batch_stride_A,
+  float const *B,
+  int ldb,
+  long long int batch_stride_B,
+  float *C,
+  int ldc,
+  long long int batch_stride_C,
+  float beta,
+  int batch_count) {
+
+  using Gemm = cutlass::gemm::device::GemmBatched<
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::ColumnMajor,
+    float, cutlass::layout::ColumnMajor
+  >;
+
+  Gemm gemm_op;
+
+  cutlass::Status status = gemm_op({
+    {m, n, k},
+    {A, lda}, 
+    batch_stride_A,
+    {B, ldb}, 
+    batch_stride_B,
+    {C, ldc}, 
+    batch_stride_C,
+    {C, ldc}, 
+    batch_stride_C,
+    {alpha, beta},
+    batch_count
+  });
+
+  if (status != cutlass::Status::kSuccess) {
+    return cudaErrorUnknown;
+  }
+
+  return cudaSuccess;
+}
+
+template<typename T> 
+cudaError_t strided_batched_gemm_nn_reference(
+  int m,
+  int n,
+  int k,
+  T alpha,
+  std::vector<T> const &A, 
+  int lda,
+  long long int batch_stride_A,
+  std::vector<T> const &B, 
+  int ldb,
+  long long int batch_stride_B,
+  std::vector<T> &C, 
+  int ldc,
+  long long int batch_stride_C,
+  T beta,
+  int batch_count) {
+  /*
+  strided batched gemm NN
+  */
+  
+  cudaError_t result = cudaSuccess;
+
+  if (A.size() < lda * k * batch_count) {
+    std::cout << "the size of A is too small" << std::endl;
+    return cudaErrorInvalidValue;
+  }
+  if (B.size() < ldb * n) {
+    std::cout << "the size of B is too small" << std::endl;
+    return cudaErrorInvalidValue;
+  }
+  if (C.size() < ldc * n * batch_count) {
+    std::cout << "the size of C is too small" << std::endl;
+    return cudaErrorInvalidValue;
+  }
+  
+  for (int batch_idx = 0; batch_idx < batch_count; batch_idx++) {
+    for (int n_idx = 0; n_idx < n; n_idx++) {
+      for (int m_idx = 0; m_idx < m; m_idx++) {
+        T accum = beta * C[batch_idx * batch_stride_C + n_idx * ldc + m_idx];
+        for (int k_idx = 0; k_idx < k; k_idx++) {
+          accum += alpha 
+            * A[batch_idx * batch_stride_A + k_idx * lda + m_idx]
+            * B[batch_idx * batch_stride_B + n_idx * ldb + k_idx];
+        }
+        C[batch_idx * batch_stride_C + n_idx * ldc + m_idx] = accum;
+      }
+    }
+  }
+
+  return result;
+}
+
+int main() {
+
+  // Arbitrary problem size
+  int const m = 520;
+  int const n = 219;
+  int const k = 129;
+  int const batch_count = 17;
+
+  // A, B are non-transpose, column major
+  int const lda = m;
+  int const ldb = k * batch_count;
+  int const ldc = m;
+
+  int const count_A = batch_count * lda * k;
+  int const count_B = ldb * n;
+  int const count_C = batch_count * ldc * n;
+
+  // the memory is batched along K dimension
+  long long int batch_stride_A = static_cast<long long int>(lda) * static_cast<long long int>(k);
+  long long int batch_stride_B = static_cast<long long int>(k);
+  long long int batch_stride_C = static_cast<long long int>(ldc) * static_cast<long long int>(n);
+
+  // alpha and beta
+  float alpha = 1.0f;
+  float beta = 2.0f;
+
+  cudaError_t result = cudaSuccess;
+
+  // allocate the host memory
+  std::vector<float> host_A(count_A);
+  std::vector<float> host_B(count_B);
+  std::vector<float> host_C(count_C);
+  std::vector<float> result_C(count_C);
+
+  // allocate the device memory
+  float *A;
+  float *B;
+  float *C;
+
+  result = cudaMalloc(&A, count_A * sizeof(float));
+  if (result != cudaSuccess) {
+    std::cerr << "cudaMalloc result = " << result << std::endl;
+    return result;
+  }
+  result = cudaMalloc(&B, count_B * sizeof(float));
+  if (result != cudaSuccess) {
+    std::cerr << "cudaMalloc result = " << result << std::endl;
+    return result;
+  }
+  result = cudaMalloc(&C, count_C * sizeof(float));
+  if (result != cudaSuccess) {
+    std::cerr << "cudaMalloc result = " << result << std::endl;
+    return result;
+  }
+
+  // Limit range to avoid floating-point errors
+  int const kRange = 8;
+
+  // fill A
+  for (int b_idx = 0; b_idx < batch_count; b_idx++) {
+    for (int col_idx = 0; col_idx < k; col_idx++) {
+      for (int row_idx = 0; row_idx < m; row_idx++) {
+        host_A[row_idx + col_idx * lda + b_idx * lda * k] = static_cast<float>((row_idx + col_idx * lda + b_idx * lda * k) % kRange);
+      }
+    }
+  }
+  // fill B
+  for (int b_idx = 0; b_idx < batch_count; b_idx++) {
+    for (int col_idx = 0; col_idx < n; col_idx++) {
+      for (int row_idx = 0; row_idx < k; row_idx++) {
+        host_B[row_idx + col_idx * ldb + b_idx * k] = static_cast<float>(((n + k * ldb + batch_count * k) - (row_idx + col_idx * ldb + b_idx * k)) % kRange);
+      }
+    }
+  }
+  // fill C
+  for (int b_idx = 0; b_idx < batch_count; b_idx++) {
+    for (int col_idx = 0; col_idx < n; col_idx++) {
+      for (int row_idx = 0; row_idx < m; row_idx++) {
+        host_C[row_idx + col_idx * ldc + b_idx * ldc * n] = 1.f;
+      }
+    }
+  }
+
+  // ref memory
+  std::vector<float> ref_A(host_A);
+  std::vector<float> ref_B(host_B);
+  std::vector<float> ref_C(host_C);
+  // copy host memory to device
+  result = cudaMemcpy(A, host_A.data(), count_A * sizeof(float), cudaMemcpyHostToDevice);
+  if (result != cudaSuccess) {
+    std::cerr << "cudaMemcpy result = " << result << std::endl;
+    return result;
+  }
+  result = cudaMemcpy(B, host_B.data(), count_B * sizeof(float), cudaMemcpyHostToDevice);
+  if (result != cudaSuccess) {
+    std::cerr << "cudaMemcpy result = " << result << std::endl;
+    return result;
+  }
+  result = cudaMemcpy(C, host_C.data(), count_C * sizeof(float), cudaMemcpyHostToDevice);
+  if (result != cudaSuccess) {
+    std::cerr << "cudaMemcpy result = " << result << std::endl;
+    return result;
+  }
+
+  // run cutlass
+  result = cutlass_strided_batched_sgemm(
+    m, n, k, alpha, A, lda, batch_stride_A, B, ldb, batch_stride_B, C, ldc, batch_stride_C,
+    beta, batch_count);
+  if (result != cudaSuccess)
+    return result;
+
+  // copy device memory to host
+  result = cudaMemcpy(result_C.data(), C, count_C * sizeof(float), cudaMemcpyDeviceToHost);
+  if (result != cudaSuccess) {
+    std::cerr << "cudaMemcpy result = " << result << std::endl;
+    return result;
+  }
+
+  //compare with reference code
+  result = strided_batched_gemm_nn_reference(m, n, k, alpha, ref_A, lda, batch_stride_A, ref_B, ldb, batch_stride_B, ref_C, ldc, batch_stride_C,
+    beta, batch_count);
+  if (result != 0)
+    return result;
+
+  // Expect bit-level accuracy for this simple example
+  if (ref_C != result_C) {
+    std::cout << "CUTLASS strided batched gemm does not run correctly" << std::endl;
+    return cudaErrorUnknown;
+  }
+
+  // free memory
+  result = cudaFree(A);
+  if (result != cudaSuccess) {
+    std::cerr << "cudaFree result = " << result << std::endl;
+    return result;
+  }
+  result = cudaFree(B);
+  if (result != cudaSuccess) {
+    std::cerr << "cudaFree result = " << result << std::endl;
+    return result;
+  }
+  result = cudaFree(C);
+  if (result != cudaSuccess) {
+    std::cerr << "cudaFree result = " << result << std::endl;
+    return result;
+  }
+
+
+  if (result == cudaSuccess) {
+    std::cout << "Passed." << std::endl;
+  }
+
+  // Exit.
+  return result == cudaSuccess ? 0 : -1;
+}

cat_files/cutlass.h

@@ -0,0 +1,175 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Basic include for CUTLASS.
+*/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define CUTLASS_UNUSED(expr) do { (void)(expr); } while (0)
+
+#if defined(_MSC_VER)
+  #define CUTLASS_NOT_IMPLEMENTED() assert(0 && __FUNCSIG__)
+#else
+  #define CUTLASS_NOT_IMPLEMENTED() assert(0 && __PRETTY_FUNCTION__)
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(__NVCC__) || (defined(__clang__) && defined(__CUDA__))
+#define CUTLASS_HOST_DEVICE __forceinline__ __device__ __host__
+#define CUTLASS_DEVICE __forceinline__ __device__
+#elif defined(__CUDACC_RTC__)
+#define CUTLASS_HOST_DEVICE __forceinline__ __device__
+#define CUTLASS_DEVICE __forceinline__ __device__
+#else
+#define CUTLASS_HOST_DEVICE inline
+#define CUTLASS_DEVICE inline
+#endif
+
+/// Status code returned by CUTLASS operations
+enum class Status {
+  kSuccess,                    ///< Operation was successful.
+  kErrorMisalignedOperand,     ///< operands fail alignment requirements.
+  kErrorInvalidDataType,       ///< DataType fails requirement.
+  kErrorInvalidLayout,         ///< Layout fails alignment requirement.
+  kErrorInvalidProblem,        ///< Specified problem size is not supported by operator.
+  kErrorNotSupported,          ///< Operation is not supported on current device.
+  kErrorWorkspaceNull,         ///< The given workspace is null when it is required to be non-null.
+  kErrorInternal,              ///< An error within CUTLASS occurred.
+  kErrorArchMismatch,          ///< CUTLASS runs on a device that it was not compiled for.
+  kErrorInsufficientDriver,    ///< CUTLASS runs with a driver that is too old.
+  kInvalid                     ///< Status is unspecified.
+};
+
+/// Convert cutlass status to status strings
+CUTLASS_HOST_DEVICE
+static char const* cutlassGetStatusString(cutlass::Status status) {
+  switch (status) {
+    case cutlass::Status::kSuccess:
+      return "Success";
+    case cutlass::Status::kErrorMisalignedOperand:
+      return "Error Misaligned Operand";
+    case cutlass::Status::kErrorInvalidDataType:
+      return "Error Invalid Data Type";
+    case cutlass::Status::kErrorInvalidLayout:
+      return "Error Invalid Layout";
+    case cutlass::Status::kErrorInvalidProblem:
+      return "Error Invalid Problem";
+    case cutlass::Status::kErrorNotSupported:
+      return "Error Not Supported";
+    case cutlass::Status::kErrorWorkspaceNull:
+      return "Error Workspace Null";
+    case cutlass::Status::kErrorInternal:
+      return "Error Internal";
+    case cutlass::Status::kErrorInsufficientDriver:
+      return "Error Insufficient Driver";
+    case cutlass::Status::kErrorArchMismatch:
+      return "Erroor Architecture Mismatch";
+    case cutlass::Status::kInvalid: break;
+  }
+
+  return "Invalid status";
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define CUTLASS_ASSERT(x) assert(x)
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// CUTLASS_PRAGMA_(UNROLL|NO_UNROLL) optimization directives for the CUDA compiler.
+#if defined(__CUDA_ARCH__)
+  #if defined(__CUDACC_RTC__) || (defined(__clang__) && defined(__CUDA__))
+    #define CUTLASS_PRAGMA_UNROLL _Pragma("unroll")
+    #define CUTLASS_PRAGMA_NO_UNROLL _Pragma("unroll 1")
+  #else
+    #define CUTLASS_PRAGMA_UNROLL #pragma unroll
+    #define CUTLASS_PRAGMA_NO_UNROLL #pragma unroll 1
+  #endif
+
+  #define CUTLASS_GEMM_LOOP CUTLASS_PRAGMA_NO_UNROLL
+
+#else
+
+    #define CUTLASS_PRAGMA_UNROLL
+    #define CUTLASS_PRAGMA_NO_UNROLL
+    #define CUTLASS_GEMM_LOOP
+
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static const int MEMORY_ACCESS_SIZE = 32;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static const int NUM_THREADS_PER_WARP = 64;
+static const int NUM_THREADS_PER_HALF_WARP = NUM_THREADS_PER_WARP / 2;
+static const int NUM_THREADS_PER_QUAD = 4;
+static const int NUM_THREADS_PER_QUAD_PAIR = NUM_THREADS_PER_QUAD * 2;
+
+#if defined(__NVCC__) || (defined(__clang__) && defined(__CUDA__))
+
+/// Computes laneId within a warp
+CUTLASS_DEVICE
+int LaneId() {
+  return __ivcorex_lane_id();
+}
+
+/// Computes SM number the thread is running on
+CUTLASS_DEVICE
+int SmId() {
+  /// TODO(Peter Han): BI compiler doesn't support sm ID
+  __asm__ __volatile__("int3");
+  return 0;
+}
+
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+

cat_files/default_gemm.h

@@ -0,0 +1,383 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+
+/*! \file
+    \brief
+      Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with
+      the appropriate threadblock-scoped epilogue.
+
+      Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are
+      accommodated by exchanging A and B operands and assuming transposed layouts. Partial
+      specializations here choose 'device::GemmTransposed' to implement this functionality.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/layout/matrix.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/arch/mma.h"
+
+#include "cutlass/epilogue/threadblock/epilogue.h"
+#include "cutlass/epilogue/thread/linear_combination.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/kernel/gemm.h"
+#include "cutlass/gemm/kernel/gemm_pipelined.h"
+#include "cutlass/gemm/threadblock/default_mma.h"
+#include "cutlass/gemm/threadblock/default_mma_core_simt.h"
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+
+#include "cutlass/epilogue/threadblock/default_epilogue_simt.h"
+#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
+#include "cutlass/transform/threadblock/predicated_tile_iterator.h"
+
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// If true, kernel is configured to support serial reduction in the
+    /// epilogue
+    bool SplitKSerial,
+    /// Operation performed by GEMM
+    typename Operator>
+struct DefaultGemm;
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for SIMT
+template <
+    /// Element type for A matrix operand
+    typename ElementA,
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// If true, kernel is configured to support serial reduction in the epilogue
+    bool SplitKSerial,
+    /// Operation performed by GEMM
+    typename Operator
+  >
+struct DefaultGemm<
+    ElementA,
+    LayoutA,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    kAlignmentB,
+    ElementC,
+    layout::RowMajor,
+    ElementAccumulator,
+    arch::OpClassSimt,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    GemmShape<1, 1, 1>,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    2,
+    SplitKSerial,
+    Operator> {
+  /// Define the threadblock-scoped matrix multiply-accumulate
+  using Mma = typename cutlass::gemm::threadblock::DefaultMma<
+      ElementA,
+      LayoutA,
+      kAlignmentA,
+      ElementB,
+      LayoutB,
+      kAlignmentB,
+      ElementAccumulator,
+      layout::RowMajor,
+      arch::OpClassSimt,
+      arch::Sm50,
+      ThreadblockShape,
+      WarpShape,
+      GemmShape<1, 1, 1>,
+      2,
+      Operator>::ThreadblockMma;
+
+  static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount;
+  static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars");
+
+  /// Define the epilogue
+  using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueSimt<
+      ThreadblockShape,
+      typename Mma::Operator,
+      EpilogueOutputOp,
+      kEpilogueElementsPerAccess
+      >::Epilogue;
+
+  /// Define the kernel-level GEMM operator.
+  using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+/// Partial specialization for SIMT DP4A
+
+template <
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentB,
+    /// Layout type for C matrix operand
+    typename LayoutC,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// If true, kernel is configured to support serial reduction in the
+    /// epilogue
+    bool SplitKSerial,
+    /// Operation performed by GEMM
+    typename Operator>
+struct DefaultGemm<int8_t, LayoutA, kAlignmentA, int8_t, LayoutB, kAlignmentB,
+                   ElementC, LayoutC, ElementAccumulator, arch::OpClassSimt,
+                   ArchTag, ThreadblockShape, WarpShape, GemmShape<1, 1, 4>,
+                   EpilogueOutputOp, ThreadblockSwizzle, 2, SplitKSerial,
+                   Operator> {
+  using InstructionShape = GemmShape<1, 1, 4>;
+  using ElementA = int8_t;
+  using ElementB = int8_t;
+
+  using OperatorClass =  arch::OpClassSimt;
+  /// Define the threadblock-scoped matrix multiply-accumulate
+  using Mma = typename cutlass::gemm::threadblock::DefaultMma<ElementA,
+      LayoutA,
+      kAlignmentA,
+      ElementB,
+      LayoutB,
+      kAlignmentB,
+      ElementAccumulator,
+      LayoutC,
+      arch::OpClassSimt,
+      arch::Sm50,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      2,
+      Operator,
+      false
+      >::ThreadblockMma;
+
+  static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount;
+  static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars");
+
+  /// Define the epilogue
+  using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueSimt<
+      ThreadblockShape,
+      typename Mma::Operator,
+      EpilogueOutputOp,
+      kEpilogueElementsPerAccess
+      >::Epilogue;
+
+  /// Define the kernel-level GEMM operator.
+  using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+};
+
+
+////////////////////////////////////////////////////////////////////////////////
+/// Partial specialization for BigIsland 1.0 tensor op architecture
+template <
+  /// Element type for A matrix operand
+  typename ElementA,
+  /// Layout type for A matrix operand
+  typename LayoutA,
+  /// Access granularity of A matrix in units of elements
+  int kAlignmentA,
+  /// Element type for B matrix operand
+  typename ElementB,
+  /// Layout type for B matrix operand
+  typename LayoutB,
+  /// Access granularity of B matrix in units of elements
+  int kAlignmentB,
+  /// Element type for C and D matrix operands
+  typename ElementC,
+  /// Element type for internal accumulation
+  typename ElementAccumulator,
+  /// Tag indicating architecture to tune for
+  typename ArchTag,
+  /// Threadblock-level tile size (concept: GemmShape)
+  typename ThreadblockShape,
+  /// Warp-level tile size (concept: GemmShape)
+  typename WarpShape,
+  /// Instrcution shape
+  typename InstructionShape,
+  /// Epilogue output operator
+  typename EpilogueOutputOp,
+  /// Threadblock-level swizzling operator
+  typename ThreadblockSwizzle,
+  /// Number of stages used in the pipelined mainloop
+  int Stages,
+  /// If true, kernel is configured to support serial reduction in the epilogue
+  bool SplitKSerial,
+  /// Operation performed by GEMM
+  typename Operator
+>
+struct DefaultGemm<
+  ElementA, LayoutA, kAlignmentA,
+  ElementB, LayoutB, kAlignmentB,
+  ElementC, layout::RowMajor,
+  ElementAccumulator,
+  arch::OpClassTensorOp,
+  ArchTag,
+  ThreadblockShape,
+  WarpShape,
+  InstructionShape,
+  EpilogueOutputOp,
+  ThreadblockSwizzle,
+  Stages,
+  SplitKSerial,
+  Operator
+> {
+
+  /// Define the threadblock-scoped matrix multiply-accumulate
+  using Mma = typename cutlass::gemm::threadblock::DefaultMma<
+    ElementA,
+    LayoutA,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    kAlignmentB,
+    ElementAccumulator,
+    layout::RowMajor,
+    arch::OpClassTensorOp,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    Stages,
+    Operator
+  >::ThreadblockMma;
+
+  static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
+
+  /// FIXME(Peter Han): Probably DefaultEpiloguesTensorOp should be used here, let's see
+  static const int kEpilougeElementsPerAccess = EpilogueOutputOp::kCount;
+
+  /// Define the epilogue
+  using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+    ThreadblockShape,
+    typename Mma::Operator,
+    EpilogueOutputOp,
+    kEpilougeElementsPerAccess
+  >::Epilogue;
+
+  /// Define the kernel-level GEMM operator.
+  using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+};
+
+
+////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass

cat_files/default_gemm_configuration.h

@@ -0,0 +1,292 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Definitions for GEMM structures
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/mma.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/epilogue/thread/linear_combination.h"
+#include "cutlass/epilogue/thread/linear_combination_clamp.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace device {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename OperatorClass,
+  typename ArchTag,
+  typename ElementA,
+  typename ElementB,
+  typename ElementC,
+  typename ElementAccumulator
+>
+struct DefaultGemmConfiguration;
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// FIXME(Peter Han): Need to update configuration according to perf results, so
+/// that could archieve good performance by default.
+
+template <
+  typename ArchTag,
+  typename ElementA,
+  typename ElementB,
+  typename ElementC,
+  typename ElementAccumulator>
+struct DefaultGemmConfiguration<
+  arch::OpClassSimt,
+  ArchTag,
+  ElementA,
+  ElementB,
+  ElementC,
+  ElementAccumulator> {
+
+  static int const kAlignmentA = 1;
+  static int const kAlignmentB = 1;
+  using ThreadblockShape = GemmShape<128, 128, 8>;
+  using WarpShape = GemmShape<64, 64, 8>;
+  using InstructionShape = GemmShape<1, 1, 1>;
+  static int const kStages = 2;
+
+  using EpilogueOutputOp = epilogue::thread::LinearCombination<
+    ElementC,
+    1,
+    ElementAccumulator,
+    ElementAccumulator
+  >;
+
+  using Operator = arch::OpMultiplyAdd;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ArchTag,
+  typename ElementC>
+struct DefaultGemmConfiguration<arch::OpClassSimt, ArchTag, int8_t, int8_t, ElementC, int32_t> {
+
+  static int const kAlignmentA = 4;
+  static int const kAlignmentB = 4;
+  using ThreadblockShape = GemmShape<128, 128, 32>;
+  using WarpShape = GemmShape<64, 64, 32>;
+  using InstructionShape = GemmShape<1, 1, 4>;
+  static int const kStages = 2;
+
+  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
+    ElementC,
+    1,
+    int32_t,
+    float
+  >;
+
+  using Operator = arch::OpMultiplyAdd;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ElementC>
+struct DefaultGemmConfiguration<
+  arch::OpClassTensorOp,
+  arch::Cu10,
+  int8_t,
+  int8_t,
+  ElementC,
+  int32_t> {
+
+  using ElementA = int8_t;
+  using ElementB = int8_t;
+  using ElementAccumulator = int32_t;
+  static int const kAlignmentA = MEMORY_ACCESS_SIZE / sizeof_bits<ElementA>::value;
+  static int const kAlignmentB = MEMORY_ACCESS_SIZE / sizeof_bits<ElementB>::value;
+
+  using ThreadblockShape = GemmShape<256, 256, 32>;
+  using WarpShape = GemmShape<64, 64, 32>;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  static int const kStages = 2;
+
+  using EpilogueOutputOp = epilogue::thread::LinearCombination<
+    ElementC,
+    MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
+    ElementAccumulator,
+    ElementAccumulator
+  >;
+
+  using Operator = arch::OpMultiplyAdd;
+};
+
+template <
+  typename ElementC>
+struct DefaultGemmConfiguration<
+  arch::OpClassTensorOp,
+  arch::Cu10,
+  uint8_t,
+  uint8_t,
+  ElementC,
+  uint32_t> {
+
+  using ElementA = uint8_t;
+  using ElementB = uint8_t;
+  using ElementAccumulator = uint32_t;
+  static int const kAlignmentA = MEMORY_ACCESS_SIZE / sizeof_bits<ElementA>::value;
+  static int const kAlignmentB = MEMORY_ACCESS_SIZE / sizeof_bits<ElementB>::value;
+
+  using ThreadblockShape = GemmShape<256, 256, 32>;
+  using WarpShape = GemmShape<64, 64, 32>;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  static int const kStages = 2;
+
+  using EpilogueOutputOp = epilogue::thread::LinearCombination<
+    ElementC,
+    MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
+    ElementAccumulator,
+    ElementAccumulator
+  >;
+
+  using Operator = arch::OpMultiplyAdd;
+};
+
+template <
+  typename ElementC>
+struct DefaultGemmConfiguration<
+  arch::OpClassTensorOp,
+  arch::Cu10,
+  half_t,
+  half_t,
+  ElementC,
+  float> {
+
+  using ElementA = half_t;
+  using ElementB = half_t;
+  using ElementAccumulator = float;
+  static int const kAlignmentA = MEMORY_ACCESS_SIZE / sizeof_bits<ElementA>::value;
+  static int const kAlignmentB = MEMORY_ACCESS_SIZE / sizeof_bits<ElementB>::value;
+
+  using ThreadblockShape = GemmShape<128, 128, 32>;
+  using WarpShape = GemmShape<32, 32, 32>;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  static int const kStages = 2;
+
+  using EpilogueOutputOp = epilogue::thread::LinearCombination<
+    ElementC,
+    MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
+    ElementAccumulator,
+    ElementAccumulator
+  >;
+
+  using Operator = arch::OpMultiplyAdd;
+};
+
+template <
+  typename ElementC>
+struct DefaultGemmConfiguration<
+  arch::OpClassTensorOp,
+  arch::Cu10,
+  bfloat16_t,
+  bfloat16_t,
+  ElementC,
+  float> {
+
+  using ElementA = bfloat16_t;
+  using ElementB = bfloat16_t;
+  using ElementAccumulator = float;
+  static int const kAlignmentA = 32 / sizeof_bits<ElementA>::value;
+  static int const kAlignmentB = 32 / sizeof_bits<ElementB>::value;
+
+  using ThreadblockShape = GemmShape<128, 128, 32>;
+  using WarpShape = GemmShape<32, 32, 32>;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  static int const kStages = 2;
+
+  using EpilogueOutputOp = epilogue::thread::LinearCombination<
+    ElementC,
+    MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
+    ElementAccumulator,
+    ElementAccumulator
+  >;
+
+  using Operator = arch::OpMultiplyAdd;
+};
+
+template <
+  typename ElementC>
+struct DefaultGemmConfiguration<
+  arch::OpClassTensorOp,
+  arch::Cu10,
+  float,
+  float,
+  ElementC,
+  float> {
+
+  using ElementA = float;
+  using ElementB = float;
+  using ElementAccumulator = float;
+  static int const kAlignmentA = 32 / sizeof_bits<ElementA>::value;
+  static int const kAlignmentB = 32 / sizeof_bits<ElementB>::value;
+
+  using ThreadblockShape = GemmShape<128, 128, 32>;
+  using WarpShape = GemmShape<32, 32, 32>;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  static int const kStages = 2;
+
+  using EpilogueOutputOp = epilogue::thread::LinearCombination<
+    ElementC,
+    MEMORY_ACCESS_SIZE / sizeof_bits<ElementC>::value,
+    ElementAccumulator,
+    ElementAccumulator
+  >;
+
+  using Operator = arch::OpMultiplyAdd;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+} // namespace device
+} // namespace gemm
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

cat_files/default_gemm_universal.h

@@ -0,0 +1,307 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief
+      Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with
+      the appropriate threadblock-scoped epilogue.
+
+      Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are
+      accommodated by exchanging A and B operands and assuming transposed layouts. Partial
+      specializations here choose 'device::GemmTransposed' to implement this functionality.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/complex.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/numeric_types.h"
+
+#include "cutlass/gemm/kernel/gemm_universal.h"
+#include "cutlass/gemm/kernel/default_gemm.h"
+#include "cutlass/gemm/kernel/default_gemm_complex.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Complex elementwise transformation on A operand
+    ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Complex elementwise transformation on B operand
+    ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator,
+    ///
+    typename Enable = void
+    >
+struct DefaultGemmUniversal;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Real-valued GEMM kernels
+//
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA,
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Layout type for C and D matrix operands
+    typename LayoutC,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator>
+struct DefaultGemmUniversal<
+  ElementA,
+  LayoutA,
+  ComplexTransform::kNone,   // transform A
+  kAlignmentA,
+  ElementB,
+  LayoutB,
+  ComplexTransform::kNone,   // transform B
+  kAlignmentB,
+  ElementC,
+  LayoutC,
+  ElementAccumulator,
+  OperatorClass,
+  ArchTag,
+  ThreadblockShape,
+  WarpShape,
+  InstructionShape,
+  EpilogueOutputOp,
+  ThreadblockSwizzle,
+  Stages,
+  Operator,
+  typename std::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
+> {
+
+  using DefaultGemmKernel = typename kernel::DefaultGemm<
+    ElementA,
+    LayoutA,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    kAlignmentB,
+    ElementC,
+    LayoutC,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    true,
+    Operator
+  >::GemmKernel;
+
+    /// Define the kernel in terms of the default kernel
+  using GemmKernel = kernel::GemmUniversal<
+    typename DefaultGemmKernel::Mma,
+    typename DefaultGemmKernel::Epilogue,
+    ThreadblockSwizzle
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Complex-valued GEMM kernels
+//
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA,
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Complex elementwise transformation on A operand
+    ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Complex elementwise transformation on B operand
+    ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Layout type for C and D matrix operands
+    typename LayoutC,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator
+  >
+struct DefaultGemmUniversal<
+  ElementA,
+  LayoutA,
+  TransformA,
+  kAlignmentA,
+  ElementB,
+  LayoutB,
+  TransformB,
+  kAlignmentB,
+  ElementC,
+  LayoutC,
+  ElementAccumulator,
+  OperatorClass,
+  ArchTag,
+  ThreadblockShape,
+  WarpShape,
+  InstructionShape,
+  EpilogueOutputOp,
+  ThreadblockSwizzle,
+  Stages,
+  Operator,
+  typename std::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
+> {
+
+  using DefaultGemmKernel = typename kernel::DefaultGemmComplex<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    TransformA,
+    TransformB,
+    Operator,
+    false
+  >::GemmKernel;
+
+  /// Define the kernel in terms of the default kernel
+  using GemmKernel = kernel::GemmUniversal<
+    typename DefaultGemmKernel::Mma,
+    typename DefaultGemmKernel::Epilogue,
+    ThreadblockSwizzle
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

cat_files/default_mma_core.h

@@ -0,0 +1,114 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data
+      layout of the global memory fragments, data types, and internal tile sizes.
+
+      Partial specializations for threadblock::Mma operations targeting TensorOp instructions.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/array.h"
+
+#include "cutlass/numeric_types.h"
+#include "cutlass/matrix_shape.h"
+
+#include "cutlass/gemm/warp/mma.h"
+#include "cutlass/gemm/threadblock/mma_pipelined.h"
+#include "cutlass/gemm/threadblock/mma_singlestage.h"
+#include "cutlass/gemm/threadblock/mma_preload.h"
+#include "cutlass/arch/cache_operation.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template defininng default matrix multiply operators inferred from threadblock tile size,
+/// global memory data layout, and target math instruction.
+template <
+    /// Shape of threadblock-scoped matrix multiply operator
+    typename Shape,
+    /// Shape of warp-level matrix multiply operator
+    typename WarpShape,
+    /// Shape of one matrix production operation (concept: GemmShape)
+    typename InstructionShape,
+    /// Element data type of A operand
+    typename ElementA,
+    /// Layout of operand A
+    typename LayoutA,
+    /// Element data type of B operand
+    typename ElementB,
+    /// Layout of operand B
+    typename LayoutB,
+    /// Data type of accumulator
+    typename ElementC,
+    /// Layout of accumulator
+    typename LayoutC,
+    /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp)
+    typename OperatorClass,
+    /// Number of stages
+    int Stages = 2,
+    /// Operation performed by MMA
+    typename Operator = cutlass::arch::OpMultiplyAdd,
+    /// Store the accumulators in row major or column major.  Row major is used
+    /// when output layout is interleaved.
+    bool AccumulatorsInRowMajor = false,
+    /// Cache operation of operand A
+    cutlass::arch::CacheOperation::Kind CacheOpA =
+        cutlass::arch::CacheOperation::Global,
+    /// Cache operation of operand B
+    cutlass::arch::CacheOperation::Kind CacheOpB =
+        cutlass::arch::CacheOperation::Global,
+    /// per-element transformation for elements of A
+    ComplexTransform TransformA = ComplexTransform::kNone,
+    /// per-element transformation for elements of B
+    ComplexTransform TransformB = ComplexTransform::kNone,
+    bool IsComplex = false // (is_complex<ElementA>::value || is_complex<ElementB>::value)
+>
+struct DefaultMmaCore;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace gemm
+} // namespace cutlass

cat_files/default_mma_core_cu10.h

@@ -0,0 +1,835 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data
+      layout of the global memory fragments, data types, and internal tile sizes.
+
+      Partial specializations for threadblock::Mma operations targeting TensorOp instructions.
+
+      Aims at TensorOp of the first generation BigIsland.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/array.h"
+
+#include "cutlass/numeric_types.h"
+#include "cutlass/matrix_shape.h"
+
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h"
+#include "cutlass/transform/threadblock/regular_tile_iterator_tensor_op.h"
+#include "cutlass/layout/tensor_op_multiplicand.h"
+#include "cutlass/layout/tensor_op_em.h"
+
+#include "cutlass/gemm/warp/mma_tensor_op_policy.h"
+#include "cutlass/gemm/warp/mma_tensor_op.h"
+#include "cutlass/gemm/warp/default_mma_tensor_op.h"
+#include "cutlass/gemm/threadblock/default_mma_core.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+///
+/// Specialization: A: row-major, B: row-major, TT
+///
+/// This uses the default warp-level operator given tile sizes
+///
+template <
+    /// Shape of threadblock-scoped matrix multiply operator (concept:
+    /// GemmShape)
+    typename Shape_,
+    /// Shape of warp-level matrix multiply operator (concept: GemmShape)
+    typename WarpShape_,
+    /// Data type of A operand
+    typename ElementA_,
+    /// Data type of B operand
+    typename ElementB_,
+    /// Data type of accumulator
+    typename ElementC_,
+    /// Layout of accumulator
+    typename LayoutC_,
+    /// Stages
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator_>
+struct DefaultMmaCore<Shape_,
+                      WarpShape_,
+                      GemmShape<16, 16, 16>,
+                      ElementA_,
+                      layout::RowMajor,
+                      ElementB_,
+                      layout::RowMajor,
+                      ElementC_,
+                      LayoutC_,
+                      arch::OpClassTensorOp,
+                      Stages,
+                      Operator_> {
+  using Shape = Shape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  using ElementA = ElementA_;
+  using LayoutA = layout::RowMajor;
+  using ElementB = ElementB_;
+  using LayoutB = layout::RowMajor;
+  using ElementC = ElementC_;
+  using LayoutC = LayoutC_;
+  using OperatorClass = arch::OpClassTensorOp;
+
+  static int const kStages = Stages;
+
+  /// Default Operator
+  using Operator = Operator_;
+
+  /// Warp thread arrangement
+  using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
+
+  /// Number of warps present
+  using WarpCount = GemmShape<
+    Shape::kM / WarpShape::kM,
+    Shape::kN / WarpShape::kN,
+    Shape::kK / WarpShape::kK
+  >;
+
+  /// Don't support split K within CTA
+  static_assert(Shape::kK == WarpShape::kK,
+    "Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(Shape::kM % WarpShape::kM) &&
+    !(Shape::kN % WarpShape::kN) &&
+    !(Shape::kK % WarpShape::kK),
+    "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(WarpShape::kM % 16) &&
+    !(WarpShape::kN % 16) &&
+    !(WarpShape::kK % 16),
+    "Threadblock-scoped GEMM should be divisible by 16."
+  );
+
+  /// Number of threads per warp
+  static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
+
+  /// Number of threads total
+  static int const kThreads = WarpCount::kCount * kWarpSize;
+
+  /// Size of a threadblock-scoped access
+  static int const kAccessSizeInBits = 32;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
+
+  //
+  // Shared memory layouts
+  //
+
+  #if BLOCK_LOAD_STORE
+    using SmemLayoutA = layout::TensorOpEm<sizeof_bits<ElementA>::value, LayoutA>;
+    using SmemLayoutB = layout::TensorOpEm<sizeof_bits<ElementB>::value, LayoutB>;
+  #else
+    using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
+    using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
+  #endif
+
+  //
+  // Iterators to write to shared memory
+  //
+
+  /// ThreadMap of iterator A
+  ///
+  using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kK, Shape::kM>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
+  >;
+
+  /// Shared memory iterator to A operand
+  using SmemIteratorA = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kM, Shape::kK>,
+    ElementA,
+    SmemLayoutA,
+    1,
+    IteratorThreadMapA
+  >;
+
+  /// Policy of iterator B
+  using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kN, Shape::kK>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
+  >;
+
+  /// Shared memory iterator to B operand
+  using SmemIteratorB = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kK, Shape::kN>,
+    ElementB,
+    SmemLayoutB,
+    0,
+    IteratorThreadMapB
+  >;
+
+  //
+  // Warp-level matrix multiply operator
+  //
+
+  // Define the warp-level tensor op
+  using Policy = gemm::warp::MmaTensorOpPolicy<
+    arch::Mma<
+      gemm::GemmShape<16, 16, 16>,
+      NUM_THREADS_PER_WARP,
+      ElementA,
+      LayoutA,
+      ElementB,
+      LayoutB,
+      ElementC,
+      layout::RowMajor,
+      arch::OpMultiplyAdd
+    >,
+    MatrixShape<1, 1>
+  >;
+
+  using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
+    WarpShape,
+    gemm::GemmShape<16, 16, 16>,
+    ElementA,
+    SmemLayoutA,
+    ElementB,
+    SmemLayoutB,
+    ElementC,
+    LayoutC,
+    arch::OpMultiplyAdd
+  >::Type;
+
+  /// Policy used to define MmaPipelined
+  using MmaPolicy = MmaPolicy<
+    MmaTensorOp,
+    MatrixShape<0, 0>,
+    MatrixShape<0, 0>,
+    WarpCount::kK
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+///
+/// Specialization: A: row-major, B: column-major, TN
+///
+/// This uses the default warp-level operator given tile sizes
+///
+template <
+    /// Shape of threadblock-scoped matrix multiply operator (concept:
+    /// GemmShape)
+    typename Shape_,
+    /// Shape of warp-level matrix multiply operator (concept: GemmShape)
+    typename WarpShape_,
+    /// Data type of A operand
+    typename ElementA_,
+    /// Data type of B operand
+    typename ElementB_,
+    /// Data type of accumulator
+    typename ElementC_,
+    /// Layout of accumulator
+    typename LayoutC_,
+    /// Stages
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator_>
+struct DefaultMmaCore<Shape_,
+                      WarpShape_,
+                      GemmShape<16, 16, 16>,
+                      ElementA_,
+                      layout::RowMajor,
+                      ElementB_,
+                      layout::ColumnMajor,
+                      ElementC_,
+                      LayoutC_,
+                      arch::OpClassTensorOp,
+                      Stages,
+                      Operator_> {
+  using Shape = Shape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  using ElementA = ElementA_;
+  using LayoutA = layout::RowMajor;
+  using ElementB = ElementB_;
+  using LayoutB = layout::ColumnMajor;
+  using ElementC = ElementC_;
+  using LayoutC = LayoutC_;
+  using OperatorClass = arch::OpClassTensorOp;
+
+  static int const kStages = Stages;
+
+  /// Default Operator
+  using Operator = Operator_;
+
+  /// Warp thread arrangement
+  using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
+
+  /// Number of warps present
+  using WarpCount = GemmShape<
+    Shape::kM / WarpShape::kM,
+    Shape::kN / WarpShape::kN,
+    Shape::kK / WarpShape::kK
+  >;
+
+  /// Don't support split K within CTA
+  static_assert(Shape::kK == WarpShape::kK,
+    "Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(Shape::kM % WarpShape::kM) &&
+    !(Shape::kN % WarpShape::kN) &&
+    !(Shape::kK % WarpShape::kK),
+    "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(WarpShape::kM % 16) &&
+    !(WarpShape::kN % 16) &&
+    !(WarpShape::kK % 16),
+    "Threadblock-scoped GEMM should be divisible by 16."
+  );
+
+  /// Number of threads per warp
+  static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
+
+  /// Number of threads total
+  static int const kThreads = WarpCount::kCount * kWarpSize;
+
+  /// Size of a threadblock-scoped access
+  static int const kAccessSizeInBits = 32;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
+
+  //
+  // Shared memory layouts
+  //
+
+  #if BLOCK_LOAD_STORE
+    using SmemLayoutA = layout::TensorOpEm<sizeof_bits<ElementA>::value, LayoutA>;
+    using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
+  #else
+    using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
+    using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
+  #endif
+
+  //
+
+  //
+  // Iterators to write to shared memory
+  //
+
+  /// ThreadMap of iterator A
+  ///
+  using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kK, Shape::kM>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
+  >;
+
+  /// Shared memory iterator to A operand
+  using SmemIteratorA = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kM, Shape::kK>,
+    ElementA,
+    SmemLayoutA,
+    1,
+    IteratorThreadMapA
+  >;
+
+  /// Policy of iterator B
+  using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kK, Shape::kN>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
+  >;
+
+  /// Shared memory iterator to B operand
+  using SmemIteratorB = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kK, Shape::kN>,
+    ElementB,
+    SmemLayoutB,
+    0,
+    IteratorThreadMapB
+  >;
+
+  //
+  // Warp-level matrix multiply operator
+  //
+
+  // Define the warp-level tensor op
+  using Policy = gemm::warp::MmaTensorOpPolicy<
+    arch::Mma<
+      gemm::GemmShape<16, 16, 16>,
+      NUM_THREADS_PER_WARP,
+      ElementA,
+      LayoutA,
+      ElementB,
+      LayoutB,
+      ElementC,
+      layout::RowMajor,
+      arch::OpMultiplyAdd
+    >,
+    MatrixShape<1, 1>
+  >;
+
+  using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
+    WarpShape,
+    gemm::GemmShape<16, 16, 16>,
+    ElementA,
+    SmemLayoutA,
+    ElementB,
+    SmemLayoutB,
+    ElementC,
+    LayoutC,
+    arch::OpMultiplyAdd
+  >::Type;
+
+  /// Policy used to define MmaPipelined
+  using MmaPolicy = MmaPolicy<
+    MmaTensorOp,
+    MatrixShape<0, 0>,
+    MatrixShape<0, 0>,
+    WarpCount::kK
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+///
+/// Specialization: A: column-major, B: row-major, NT
+///
+/// This uses the default warp-level operator given tile sizes
+///
+template <
+    /// Shape of threadblock-scoped matrix multiply operator (concept:
+    /// GemmShape)
+    typename Shape_,
+    /// Shape of warp-level matrix multiply operator (concept: GemmShape)
+    typename WarpShape_,
+    /// Data type of A operand
+    typename ElementA_,
+    /// Data type of B operand
+    typename ElementB_,
+    /// Data type of accumulator
+    typename ElementC_,
+    /// Layout of accumulator
+    typename LayoutC_,
+    /// Stages
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator_>
+struct DefaultMmaCore<Shape_,
+                      WarpShape_,
+                      GemmShape<16, 16, 16>,
+                      ElementA_,
+                      layout::ColumnMajor,
+                      ElementB_,
+                      layout::RowMajor,
+                      ElementC_,
+                      LayoutC_,
+                      arch::OpClassTensorOp,
+                      Stages,
+                      Operator_> {
+  using Shape = Shape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  using ElementA = ElementA_;
+  using LayoutA = layout::ColumnMajor;
+  using ElementB = ElementB_;
+  using LayoutB = layout::RowMajor;
+  using ElementC = ElementC_;
+  using LayoutC = LayoutC_;
+  using OperatorClass = arch::OpClassTensorOp;
+
+  static int const kStages = Stages;
+
+  /// Default Operator
+  using Operator = Operator_;
+
+  /// Warp thread arrangement
+  using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
+
+  /// Number of warps present
+  using WarpCount = GemmShape<
+    Shape::kM / WarpShape::kM,
+    Shape::kN / WarpShape::kN,
+    Shape::kK / WarpShape::kK
+  >;
+
+  /// Don't support split K within CTA
+  static_assert(Shape::kK == WarpShape::kK,
+    "Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(Shape::kM % WarpShape::kM) &&
+    !(Shape::kN % WarpShape::kN) &&
+    !(Shape::kK % WarpShape::kK),
+    "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(WarpShape::kM % 16) &&
+    !(WarpShape::kN % 16) &&
+    !(WarpShape::kK % 16),
+    "Threadblock-scoped GEMM should be divisible by 16."
+  );
+
+  /// Number of threads per warp
+  static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
+
+  /// Number of threads total
+  static int const kThreads = WarpCount::kCount * kWarpSize;
+
+  /// Size of a threadblock-scoped access
+  static int const kAccessSizeInBits = 32;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
+
+  //
+  // Shared memory layouts
+  //
+
+  #if BLOCK_LOAD_STORE
+    using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
+    using SmemLayoutB = layout::TensorOpEm<sizeof_bits<ElementB>::value, LayoutB>;
+  #else
+    using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
+    using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
+  #endif
+
+  //
+  // Iterators to write to shared memory
+  //
+
+  /// ThreadMap of iterator A
+  ///
+  using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kM, Shape::kK>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
+  >;
+
+  /// Shared memory iterator to A operand
+  using SmemIteratorA = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kM, Shape::kK>,
+    ElementA,
+    SmemLayoutA,
+    1,
+    IteratorThreadMapA
+  >;
+
+  /// Policy of iterator B
+  using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kN, Shape::kK>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
+  >;
+
+  /// Shared memory iterator to B operand
+  using SmemIteratorB = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kK, Shape::kN>,
+    ElementB,
+    SmemLayoutB,
+    0,
+    IteratorThreadMapB
+  >;
+
+  //
+  // Warp-level matrix multiply operator
+  //
+
+  // Define the warp-level tensor op
+  using Policy = gemm::warp::MmaTensorOpPolicy<
+    arch::Mma<
+      gemm::GemmShape<16, 16, 16>,
+      NUM_THREADS_PER_WARP,
+      ElementA,
+      LayoutA,
+      ElementB,
+      LayoutB,
+      ElementC,
+      layout::RowMajor,
+      arch::OpMultiplyAdd
+    >,
+    MatrixShape<1, 1>
+  >;
+
+  using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
+    WarpShape,
+    gemm::GemmShape<16, 16, 16>,
+    ElementA,
+    SmemLayoutA,
+    ElementB,
+    SmemLayoutB,
+    ElementC,
+    LayoutC,
+    arch::OpMultiplyAdd
+  >::Type;
+
+  /// Policy used to define MmaPipelined
+  using MmaPolicy = MmaPolicy<
+    MmaTensorOp,
+    MatrixShape<0, 0>,
+    MatrixShape<0, 0>,
+    WarpCount::kK
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+///
+/// Specialization: A: column-major, B: column-major, NN
+///
+/// This uses the default warp-level operator given tile sizes
+///
+template <
+    /// Shape of threadblock-scoped matrix multiply operator (concept:
+    /// GemmShape)
+    typename Shape_,
+    /// Shape of warp-level matrix multiply operator (concept: GemmShape)
+    typename WarpShape_,
+    /// Data type of A operand
+    typename ElementA_,
+    /// Data type of B operand
+    typename ElementB_,
+    /// Data type of accumulator
+    typename ElementC_,
+    /// Layout of accumulator
+    typename LayoutC_,
+    /// Stages
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator_>
+struct DefaultMmaCore<Shape_,
+                      WarpShape_,
+                      GemmShape<16, 16, 16>,
+                      ElementA_,
+                      layout::ColumnMajor,
+                      ElementB_,
+                      layout::ColumnMajor,
+                      ElementC_,
+                      LayoutC_,
+                      arch::OpClassTensorOp,
+                      Stages,
+                      Operator_> {
+  using Shape = Shape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = GemmShape<16, 16, 16>;
+  using ElementA = ElementA_;
+  using LayoutA = layout::ColumnMajor;
+  using ElementB = ElementB_;
+  using LayoutB = layout::ColumnMajor;
+  using ElementC = ElementC_;
+  using LayoutC = LayoutC_;
+  using OperatorClass = arch::OpClassTensorOp;
+
+  static int const kStages = Stages;
+
+  /// Default Operator
+  using Operator = Operator_;
+
+  /// Warp thread arrangement
+  using WarpThreadArrangement = layout::PitchLinearShape<16, 4>;
+
+  /// Number of warps present
+  using WarpCount = GemmShape<
+    Shape::kM / WarpShape::kM,
+    Shape::kN / WarpShape::kN,
+    Shape::kK / WarpShape::kK
+  >;
+
+  /// Don't support split K within CTA
+  static_assert(Shape::kK == WarpShape::kK,
+    "Threadblock-scoped GEMM shape K should equal warp-scoped GEMM shape K"
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(Shape::kM % WarpShape::kM) &&
+    !(Shape::kN % WarpShape::kN) &&
+    !(Shape::kK % WarpShape::kK),
+    "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."
+  );
+
+  // Divisibility requirements
+  static_assert(
+    !(WarpShape::kM % 16) &&
+    !(WarpShape::kN % 16) &&
+    !(WarpShape::kK % 16),
+    "Threadblock-scoped GEMM should be divisible by 16."
+  );
+
+  /// Number of threads per warp
+  static int const kWarpSize = warp::WarpSize<arch::OpClassTensorOp>::value;
+
+  /// Number of threads total
+  static int const kThreads = WarpCount::kCount * kWarpSize;
+
+  /// Size of a threadblock-scoped access
+  static int const kAccessSizeInBits = 32;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessA = kAccessSizeInBits / sizeof_bits<ElementA>::value;
+
+  /// Number of A elemnts per access
+  static int const kElementsPerAccessB = kAccessSizeInBits / sizeof_bits<ElementB>::value;
+
+  //
+  // Shared memory layouts
+  //
+  using SmemLayoutA = layout::TensorOpMultiplicand<sizeof_bits<ElementA>::value, LayoutA>;
+  using SmemLayoutB = layout::TensorOpMultiplicand<sizeof_bits<ElementB>::value, LayoutB>;
+
+  //
+
+  //
+  // Iterators to write to shared memory
+  //
+
+  /// ThreadMap of iterator A
+  ///
+  using IteratorThreadMapA = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kM, Shape::kK>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessA, kElementsPerAccessA>
+  >;
+
+  /// Shared memory iterator to A operand
+  using SmemIteratorA = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kM, Shape::kK>,
+    ElementA,
+    SmemLayoutA,
+    1,
+    IteratorThreadMapA
+  >;
+
+  /// Policy of iterator B
+  using IteratorThreadMapB = transform::PitchLinear2DThreadTileWarpRakedThreadMap<
+    layout::PitchLinearShape<Shape::kK, Shape::kN>,
+    kThreads,
+    WarpThreadArrangement,
+    layout::PitchLinearShape<kElementsPerAccessB, kElementsPerAccessB>
+  >;
+
+  /// Shared memory iterator to B operand
+  using SmemIteratorB = transform::threadblock::RegularTileIterator<
+    MatrixShape<Shape::kK, Shape::kN>,
+    ElementB,
+    SmemLayoutB,
+    0,
+    IteratorThreadMapB
+  >;
+
+  //
+  // Warp-level matrix multiply operator
+  //
+
+  // Define the warp-level tensor op
+  using Policy = gemm::warp::MmaTensorOpPolicy<
+    arch::Mma<
+      gemm::GemmShape<16, 16, 16>,
+      NUM_THREADS_PER_WARP,
+      ElementA,
+      LayoutA,
+      ElementB,
+      LayoutB,
+      ElementC,
+      layout::RowMajor,
+      arch::OpMultiplyAdd
+    >,
+    MatrixShape<1, 1>
+  >;
+
+  using MmaTensorOp = typename gemm::warp::DefaultMmaTensorOp<
+    WarpShape,
+    gemm::GemmShape<16, 16, 16>,
+    ElementA,
+    SmemLayoutA,
+    ElementB,
+    SmemLayoutB,
+    ElementC,
+    LayoutC,
+    arch::OpMultiplyAdd
+  >::Type;
+
+  /// Policy used to define MmaPipelined
+  using MmaPolicy = MmaPolicy<
+    MmaTensorOp,
+    MatrixShape<0, 0>,
+    MatrixShape<0, 0>,
+    WarpCount::kK
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace gemm
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

cat_files/default_mma_tensor_op.h

@@ -0,0 +1,148 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/warp/mma_tensor_op.h"
+
+namespace cutlass {
+namespace gemm {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Size of the Gemm problem - concept: gemm::GemmShape<>
+    typename WarpShape_,
+    /// Shape of one matrix production operation (concept: GemmShape)
+    typename InstructionShape_,
+    /// Data type of A elements
+    typename ElementA_,
+    /// Layout of A matrix (concept: MatrixLayout)
+    typename LayoutA_,
+    /// Data type of B elements
+    typename ElementB_,
+    /// Layout of B matrix (concept: MatrixLayout)
+    typename LayoutB_,
+    /// Element type of C matrix
+    typename ElementC_,
+    /// Layout of C matrix (concept: MatrixLayout)
+    typename LayoutC_,
+    /// Operator describing the tensor operation
+    typename Operator_ = arch::OpMultiplyAdd,
+    /// Number of partitions along K dimension
+    int PartitionsK = 1,
+    /// Store the accumulators in row major or column major.
+    bool AccumulatorsInRowMajor = true>
+struct DefaultMmaTensorOp;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for m-by-n-by-kgroup
+template <
+    /// Shape of one matrix production operation (concept: GemmShape)
+    typename WarpShape_,
+    /// Data type of A elements
+    typename ElementA,
+    /// Layout of A matrix (concept: MatrixLayout)
+    typename LayoutA,
+    /// Data type of B elements
+    typename ElementB,
+    /// Layout of B matrix (concept: MatrixLayout)
+    typename LayoutB,
+    /// Element type of C matrix
+    typename ElementC,
+    /// Layout of C matrix (concept: MatrixLayout)
+    typename LayoutC,
+    /// Number of partitions along K dimension
+    int PartitionsK,
+    /// Store the accumulators in row major or column major.
+    bool AccumulatorsInRowMajor>
+struct DefaultMmaTensorOp<
+    WarpShape_,
+    GemmShape<16, 16, 16>,
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    arch::OpMultiplyAdd,
+    PartitionsK,
+    AccumulatorsInRowMajor> {
+
+  /// Warp shape
+  using Shape = WarpShape_;
+
+  using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<
+      cutlass::arch::Mma<GemmShape<16, 16, 16>,
+                         64,
+                         ElementA,
+                         LayoutA,
+                         ElementB,
+                         LayoutB,
+                         ElementC,
+                         LayoutC,
+                         arch::OpMultiplyAdd>,
+      cutlass::MatrixShape<1, 1> >;
+
+  // Define the warp-level tensor op
+  using Type = cutlass::gemm::warp::MmaTensorOp<
+      WarpShape_,
+      ElementA,
+      LayoutA,
+      ElementB,
+      LayoutB,
+      ElementC,
+      LayoutC,
+      Policy,
+      PartitionsK,
+      AccumulatorsInRowMajor>;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace gemm
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+

cat_files/gemm_batched.h

@@ -0,0 +1,726 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/device_kernel.h"
+
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+#include "cutlass/gemm/kernel/gemm_batched.h"
+
+#include "cutlass/gemm/kernel/default_gemm.h"
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace device {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/*! Gemm device-level operator. This is an interface to efficient CUTLASS GEMM kernels that may
+  be invoked from host code.
+
+  The contributions of this class are:
+
+    1. At compile time, it maps data types and high-level structural parameters onto
+       specific CUTLASS components.
+
+    2. At runtime, it maps logical arguments to GEMM problems to kernel parameters.
+
+    3. At runtime, it launches kernels on the device.
+
+  The intent is to provide a convenient mechanism for interacting with most plausible GEMM
+  configurations for each supported architecture. Consequently, not all parameters are exposed
+  to the top-level interface. Rather, sensible defaults at each level of the CUTLASS hierarchy
+  are selected to tradeoff simplicity of the interface with flexibility. We expect
+  most configurations to be specified at this level. Applications with more exotic requirements
+  may construct their kernels of interest using CUTLASS components at the threadblock, warp,
+  and thread levels of abstraction.
+
+  CUTLASS exposes computations using the functor design pattern in which objects compose some
+  internal state with an overloaded function call operator. This enables decoupling of
+  initialization from execution, possibly reducing overhead during steady state phases of
+  application execution.
+
+  CUTLASS device-level operators expose an Arguments structure encompassing each logical
+  input to the computation. This is distinct from the kernel-level Params structure pattern
+  which contains application-specific precomputed state needed by the device code.
+
+  Example of a CUTLASS GEMM operator implementing the functionality of cuBLAS's SGEMM NN
+  is as follows:
+
+    //
+    // Instantiate the CUTLASS GEMM operator.
+    //
+
+    cutlass::gemm::device::Gemm<
+      float,
+      cutlass::layout::ColumnMajor,
+      float,
+      cutlass::layout::ColumnMajor,
+      float,
+      cutlass::layout::ColumnMajor
+    > gemm_op;
+
+    //
+    // Launch the GEMM operation on the device
+    //
+
+    cutlass::Status status = gemm_op({
+      {m, n, k},                          // GemmCoord problem_size,
+      {A, lda},                           // TensorRef<float, layout::ColumnMajor> ref_A,
+      {B, ldb},                           // TensorRef<float, layout::ColumnMajor> ref_B,
+      {C, ldc},                           // TensorRef<float, layout::ColumnMajor> ref_C,
+      {D, ldd},                           // TensorRef<float, layout::ColumnMajor> ref_D,
+      {alpha, beta}                       // EpilogueOutputOp::Params epilogue_op_params
+    });
+
+
+  A simplified view of the template is listed below.
+
+    template <
+      /// Element type for A matrix operand
+      typename ElementA,
+
+      /// Layout type for A matrix operand
+      typename LayoutA,
+
+      /// Element type for B matrix operand
+      typename ElementB,
+
+      /// Layout type for B matrix operand
+      typename LayoutB,
+
+      /// Element type for C and D matrix operands
+      typename ElementC,
+
+      /// Layout type for C and D matrix operands
+      typename LayoutC,
+
+      /// Element type for internal accumulation
+      typename ElementAccumulator,
+
+      /// Operator class tag
+      typename OperatorClass,
+
+      /// Tag indicating architecture to tune for.  This is the minimum SM that
+      /// supports the intended feature. The device kernel can be built
+      /// targeting any SM larger than this number.
+      typename ArchTag,
+
+      /// Threadblock-level tile size (concept: GemmShape)
+      typename ThreadblockShape,
+
+      /// Warp-level tile size (concept: GemmShape)
+      typename WarpShape,
+
+      /// Warp-level tile size (concept: GemmShape)
+      typename InstructionShape,
+
+      /// Epilogue output operator
+      typename EpilogueOutputOp,
+
+      /// Threadblock-level swizzling operator
+      typename ThreadblockSwizzle,
+
+      /// Number of stages used in the pipelined mainloop
+      int Stages
+    >
+    class Gemm;
+*/
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_ = ElementC_,
+    /// Operator class tag
+    typename OperatorClass_ = arch::OpClassSimt,
+    /// Tag indicating architecture to tune for
+    typename ArchTag_ = arch::Sm61,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_ = threadblock::GemmBatchedIdentityThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kStages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::Operator
+>
+class GemmBatched {
+ public:
+
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = LayoutC_;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  using Operator = Operator_;
+
+  /// Define the kernel
+  using DefaultGemmKernel = typename kernel::DefaultGemm<
+    ElementA,
+    LayoutA,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    kAlignmentB,
+    ElementC,
+    LayoutC,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    kStages,
+    false,
+    Operator
+  >::GemmKernel;
+
+  using GemmKernel = kernel::GemmBatched<typename DefaultGemmKernel::Mma, typename DefaultGemmKernel::Epilogue, ThreadblockSwizzle>;
+
+  /// Argument structure
+  struct Arguments {
+
+    //
+    // Data members
+    //
+
+    GemmCoord problem_size;
+    TensorRef<ElementA const, LayoutA> ref_A;
+    int64_t stride_A;
+    TensorRef<ElementB const, LayoutB> ref_B;
+    int64_t stride_B;
+    TensorRef<ElementC const, LayoutC> ref_C;
+    int64_t stride_C;
+    TensorRef<ElementC, LayoutC> ref_D;
+    int64_t stride_D;
+    typename EpilogueOutputOp::Params epilogue;
+    int batch_count;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments() { }
+
+    /// Constructs an Arguments structure
+    CUTLASS_HOST_DEVICE
+    Arguments(
+      GemmCoord problem_size_,
+      TensorRef<ElementA const, LayoutA> ref_A_,
+      int64_t stride_A_,
+      TensorRef<ElementB const, LayoutB> ref_B_,
+      int64_t stride_B_,
+      TensorRef<ElementC const, LayoutC> ref_C_,
+      int64_t stride_C_,
+      TensorRef<ElementC, LayoutC> ref_D_,
+      int64_t stride_D_,
+      typename EpilogueOutputOp::Params epilogue_,
+      int batch_count_
+    ):
+      problem_size(problem_size_),
+      ref_A(ref_A_),
+      stride_A(stride_A_),
+      ref_B(ref_B_),
+      stride_B(stride_B_),
+      ref_C(ref_C_),
+      stride_C(stride_C_),
+      ref_D(ref_D_),
+      stride_D(stride_D_),
+      epilogue(epilogue_),
+      batch_count(batch_count_) { }
+  };
+
+private:
+
+  /// Kernel parameters object
+  typename GemmKernel::Params params_;
+
+public:
+
+  /// Constructs the GEMM.
+  GemmBatched() { }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+
+    if (!TensorRef_aligned(args.ref_A, kAlignmentA) || (args.stride_A % kAlignmentA)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(args.ref_B, kAlignmentB) || (args.stride_B % kAlignmentB)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(args.ref_C, kAlignmentC) || (args.stride_C % kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(args.ref_D, kAlignmentC) || (args.stride_D % kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if ((args.problem_size.m() % kAlignmentA) || (args.problem_size.k() % kAlignmentA) ||
+      (args.problem_size.n() % kAlignmentB) || (args.problem_size.k() % kAlignmentB) ||
+      (args.problem_size.m() % kAlignmentC) || (args.problem_size.n() % kAlignmentC)) {
+
+      return Status::kErrorMisalignedOperand;
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+    return 0;
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
+
+    // Determine grid shape
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(
+      args.problem_size,
+      {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+      args.batch_count);
+
+    // Initialize the Params structure
+    params_ = typename GemmKernel::Params{
+      args.problem_size,
+      grid_shape,
+      args.ref_A.non_const_ref(),
+      args.stride_A,
+      args.ref_B.non_const_ref(),
+      args.stride_B,
+      args.ref_C.non_const_ref(),
+      args.stride_C,
+      args.ref_D,
+      args.stride_D,
+      args.epilogue,
+      args.batch_count
+    };
+
+    return Status::kSuccess;
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+
+    params_.ref_A.reset(args.ref_A.non_const_ref().data());
+    params_.ref_B.reset(args.ref_B.non_const_ref().data());
+    params_.ref_C.reset(args.ref_C.non_const_ref().data());
+    params_.ref_D.reset(args.ref_D.data());
+
+    return Status::kSuccess;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+
+    ThreadblockSwizzle threadblock_swizzle;
+
+    dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
+    // XXX(Peter Han): prealod needs double warps in z direction
+    dim3 block(GemmKernel::kThreadCount, 1, kStages ? 1 : 2);
+
+    cudaError_t result;
+
+    int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
+    /// cudaFuncSetAttribute isn't supported under CUDA-8.0
+    // if (smem_size >= (48 << 10)) {
+    //   result = cudaFuncSetAttribute(Kernel<GemmKernel>,
+    //                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
+    //                                 smem_size);
+
+    //   if (result != cudaSuccess) {
+    //     return Status::kErrorInternal;
+    //   }
+
+    //   result = cudaFuncSetAttribute(
+    //       Kernel<GemmKernel>,
+    //       cudaFuncAttributePreferredSharedMemoryCarveout, 100);
+
+    //   if (result != cudaSuccess) {
+    //     return Status::kErrorInternal;
+    //   }
+    // }
+
+    cutlass::Kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params_);
+
+    result = cudaGetLastError();
+
+    return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(
+    Arguments const &args,
+    void *workspace = nullptr,
+    cudaStream_t stream = nullptr) {
+
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Parital specialization for column-major output exchanges problem size and operand.
+template <
+  /// Element type for A matrix operand
+  typename ElementA_,
+  /// Layout type for A matrix operand
+  typename LayoutA_,
+  /// Element type for B matrix operand
+  typename ElementB_,
+  /// Layout type for B matrix operand
+  typename LayoutB_,
+  /// Element type for C and D matrix operands
+  typename ElementC_,
+  /// Element type for internal accumulation
+  typename ElementAccumulator_,
+  /// Operator class tag
+  typename OperatorClass_,
+  /// Tag indicating architecture to tune for
+  typename ArchTag_,
+  /// Threadblock-level tile size (concept: GemmShape)
+  typename ThreadblockShape_,
+  /// Warp-level tile size (concept: GemmShape)
+  typename WarpShape_,
+  /// Warp-level tile size (concept: GemmShape)
+  typename InstructionShape_,
+  /// Epilogue output operator
+  typename EpilogueOutputOp_,
+  /// Threadblock-level swizzling operator
+  typename ThreadblockSwizzle_,
+  /// Number of stages used in the pipelined mainloop
+  int Stages,
+  /// Access granularity of A matrix in units of elements
+  int AlignmentA,
+  /// Access granularity of B matrix in units of elements
+  int AlignmentB,
+  typename Operator_
+>
+class GemmBatched<
+  ElementA_,
+  LayoutA_,
+  ElementB_,
+  LayoutB_,
+  ElementC_,
+  layout::ColumnMajor,
+  ElementAccumulator_,
+  OperatorClass_,
+  ArchTag_,
+  ThreadblockShape_,
+  WarpShape_,
+  InstructionShape_,
+  EpilogueOutputOp_,
+  ThreadblockSwizzle_,
+  Stages,
+  AlignmentA,
+  AlignmentB,
+  Operator_
+> {
+public:
+
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = layout::ColumnMajor;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static int const kStages = Stages;
+
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  static bool const kSplitKSerial = false;
+
+  //
+  using UnderlyingOperator = GemmBatched<
+    ElementB,
+    typename layout::LayoutTranspose<LayoutB>::type,
+    ElementA,
+    typename layout::LayoutTranspose<LayoutA>::type,
+    ElementC,
+    layout::RowMajor,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    kAlignmentB,
+    kAlignmentA
+  >;
+
+  using UnderlyingArguments = typename UnderlyingOperator::Arguments;
+  using GemmKernel = typename UnderlyingOperator::GemmKernel;
+
+  /// Argument structure
+  struct Arguments {
+
+    //
+    // Data members
+    //
+
+    GemmCoord problem_size;
+    TensorRef<ElementA const, LayoutA> ref_A;
+    int64_t stride_A;
+    TensorRef<ElementB const, LayoutB> ref_B;
+    int64_t stride_B;
+    TensorRef<ElementC const, LayoutC> ref_C;
+    int64_t stride_C;
+    TensorRef<ElementC, LayoutC> ref_D;
+    int64_t stride_D;
+    typename EpilogueOutputOp::Params epilogue;
+    int batch_count;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments() { }
+
+    /// Constructs an Arguments structure
+    CUTLASS_HOST_DEVICE
+    Arguments(
+      GemmCoord problem_size_,
+      TensorRef<ElementA const, LayoutA> ref_A_,
+      int64_t stride_A_,
+      TensorRef<ElementB const, LayoutB> ref_B_,
+      int64_t stride_B_,
+      TensorRef<ElementC const, LayoutC> ref_C_,
+      int64_t stride_C_,
+      TensorRef<ElementC, LayoutC> ref_D_,
+      int64_t stride_D_,
+      typename EpilogueOutputOp::Params epilogue_,
+      int batch_count_
+    ):
+      problem_size(problem_size_),
+      ref_A(ref_A_),
+      stride_A(stride_A_),
+      ref_B(ref_B_),
+      stride_B(stride_B_),
+      ref_C(ref_C_),
+      stride_C(stride_C_),
+      ref_D(ref_D_),
+      stride_D(stride_D_),
+      epilogue(epilogue_),
+      batch_count(batch_count_) { }
+  };
+
+private:
+
+  UnderlyingOperator underlying_operator_;
+
+public:
+
+  /// Constructs the GEMM.
+  GemmBatched() { }
+
+  /// Helper to construct a transposed equivalent for the underying GEMM operator
+  static UnderlyingArguments to_underlying_arguments(Arguments const &args) {
+    return UnderlyingArguments(
+      {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()},
+      {args.ref_B.data(), args.ref_B.stride(0)},
+      args.stride_B,
+      {args.ref_A.data(), args.ref_A.stride(0)},
+      args.stride_A,
+      {args.ref_C.data(), args.ref_C.stride(0)},
+      args.stride_C,
+      {args.ref_D.data(), args.ref_D.stride(0)},
+      args.stride_D,
+      args.epilogue,
+      args.batch_count
+    );
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+
+    return UnderlyingOperator::can_implement(to_underlying_arguments(args));
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+
+    return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.initialize(to_underlying_arguments(args), workspace);
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+
+    return underlying_operator_.update(to_underlying_arguments(args), workspace);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(
+    Arguments const &args,
+    void *workspace = nullptr,
+    cudaStream_t stream = nullptr) {
+
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace gemm
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

cat_files/gemm_batched_full.h

@@ -0,0 +1,726 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/device_kernel.h"
+
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+#include "cutlass/gemm/kernel/gemm_batched.h"
+
+#include "cutlass/gemm/kernel/default_gemm.h"
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace device {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/*! Gemm device-level operator. This is an interface to efficient CUTLASS GEMM kernels that may
+  be invoked from host code.
+
+  The contributions of this class are:
+
+    1. At compile time, it maps data types and high-level structural parameters onto
+       specific CUTLASS components.
+
+    2. At runtime, it maps logical arguments to GEMM problems to kernel parameters.
+
+    3. At runtime, it launches kernels on the device.
+
+  The intent is to provide a convenient mechanism for interacting with most plausible GEMM
+  configurations for each supported architecture. Consequently, not all parameters are exposed
+  to the top-level interface. Rather, sensible defaults at each level of the CUTLASS hierarchy
+  are selected to tradeoff simplicity of the interface with flexibility. We expect
+  most configurations to be specified at this level. Applications with more exotic requirements
+  may construct their kernels of interest using CUTLASS components at the threadblock, warp,
+  and thread levels of abstraction.
+
+  CUTLASS exposes computations using the functor design pattern in which objects compose some
+  internal state with an overloaded function call operator. This enables decoupling of
+  initialization from execution, possibly reducing overhead during steady state phases of
+  application execution.
+
+  CUTLASS device-level operators expose an Arguments structure encompassing each logical
+  input to the computation. This is distinct from the kernel-level Params structure pattern
+  which contains application-specific precomputed state needed by the device code.
+
+  Example of a CUTLASS GEMM operator implementing the functionality of cuBLAS's SGEMM NN
+  is as follows:
+
+    //
+    // Instantiate the CUTLASS GEMM operator.
+    //
+
+    cutlass::gemm::device::Gemm<
+      float,
+      cutlass::layout::ColumnMajor,
+      float,
+      cutlass::layout::ColumnMajor,
+      float,
+      cutlass::layout::ColumnMajor
+    > gemm_op;
+
+    //
+    // Launch the GEMM operation on the device
+    //
+
+    cutlass::Status status = gemm_op({
+      {m, n, k},                          // GemmCoord problem_size,
+      {A, lda},                           // TensorRef<float, layout::ColumnMajor> ref_A,
+      {B, ldb},                           // TensorRef<float, layout::ColumnMajor> ref_B,
+      {C, ldc},                           // TensorRef<float, layout::ColumnMajor> ref_C,
+      {D, ldd},                           // TensorRef<float, layout::ColumnMajor> ref_D,
+      {alpha, beta}                       // EpilogueOutputOp::Params epilogue_op_params
+    });
+
+
+  A simplified view of the template is listed below.
+
+    template <
+      /// Element type for A matrix operand
+      typename ElementA,
+
+      /// Layout type for A matrix operand
+      typename LayoutA,
+
+      /// Element type for B matrix operand
+      typename ElementB,
+
+      /// Layout type for B matrix operand
+      typename LayoutB,
+
+      /// Element type for C and D matrix operands
+      typename ElementC,
+
+      /// Layout type for C and D matrix operands
+      typename LayoutC,
+
+      /// Element type for internal accumulation
+      typename ElementAccumulator,
+
+      /// Operator class tag
+      typename OperatorClass,
+
+      /// Tag indicating architecture to tune for.  This is the minimum SM that
+      /// supports the intended feature. The device kernel can be built
+      /// targeting any SM larger than this number.
+      typename ArchTag,
+
+      /// Threadblock-level tile size (concept: GemmShape)
+      typename ThreadblockShape,
+
+      /// Warp-level tile size (concept: GemmShape)
+      typename WarpShape,
+
+      /// Warp-level tile size (concept: GemmShape)
+      typename InstructionShape,
+
+      /// Epilogue output operator
+      typename EpilogueOutputOp,
+
+      /// Threadblock-level swizzling operator
+      typename ThreadblockSwizzle,
+
+      /// Number of stages used in the pipelined mainloop
+      int Stages
+    >
+    class Gemm;
+*/
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_ = ElementC_,
+    /// Operator class tag
+    typename OperatorClass_ = arch::OpClassSimt,
+    /// Tag indicating architecture to tune for
+    typename ArchTag_ = arch::Sm61,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_ = threadblock::GemmBatchedIdentityThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kStages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::Operator
+>
+class GemmBatched {
+ public:
+
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = LayoutC_;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  using Operator = Operator_;
+
+  /// Define the kernel
+  using DefaultGemmKernel = typename kernel::DefaultGemm<
+    ElementA,
+    LayoutA,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    kAlignmentB,
+    ElementC,
+    LayoutC,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    kStages,
+    false,
+    Operator
+  >::GemmKernel;
+
+  using GemmKernel = kernel::GemmBatched<typename DefaultGemmKernel::Mma, typename DefaultGemmKernel::Epilogue, ThreadblockSwizzle>;
+
+  /// Argument structure
+  struct Arguments {
+
+    //
+    // Data members
+    //
+
+    GemmCoord problem_size;
+    TensorRef<ElementA const, LayoutA> ref_A;
+    int64_t stride_A;
+    TensorRef<ElementB const, LayoutB> ref_B;
+    int64_t stride_B;
+    TensorRef<ElementC const, LayoutC> ref_C;
+    int64_t stride_C;
+    TensorRef<ElementC, LayoutC> ref_D;
+    int64_t stride_D;
+    typename EpilogueOutputOp::Params epilogue;
+    int batch_count;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments() { }
+
+    /// Constructs an Arguments structure
+    CUTLASS_HOST_DEVICE
+    Arguments(
+      GemmCoord problem_size_,
+      TensorRef<ElementA const, LayoutA> ref_A_,
+      int64_t stride_A_,
+      TensorRef<ElementB const, LayoutB> ref_B_,
+      int64_t stride_B_,
+      TensorRef<ElementC const, LayoutC> ref_C_,
+      int64_t stride_C_,
+      TensorRef<ElementC, LayoutC> ref_D_,
+      int64_t stride_D_,
+      typename EpilogueOutputOp::Params epilogue_,
+      int batch_count_
+    ):
+      problem_size(problem_size_),
+      ref_A(ref_A_),
+      stride_A(stride_A_),
+      ref_B(ref_B_),
+      stride_B(stride_B_),
+      ref_C(ref_C_),
+      stride_C(stride_C_),
+      ref_D(ref_D_),
+      stride_D(stride_D_),
+      epilogue(epilogue_),
+      batch_count(batch_count_) { }
+  };
+
+private:
+
+  /// Kernel parameters object
+  typename GemmKernel::Params params_;
+
+public:
+
+  /// Constructs the GEMM.
+  GemmBatched() { }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+
+    if (!TensorRef_aligned(args.ref_A, kAlignmentA) || (args.stride_A % kAlignmentA)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(args.ref_B, kAlignmentB) || (args.stride_B % kAlignmentB)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(args.ref_C, kAlignmentC) || (args.stride_C % kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(args.ref_D, kAlignmentC) || (args.stride_D % kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if ((args.problem_size.m() % kAlignmentA) || (args.problem_size.k() % kAlignmentA) ||
+      (args.problem_size.n() % kAlignmentB) || (args.problem_size.k() % kAlignmentB) ||
+      (args.problem_size.m() % kAlignmentC) || (args.problem_size.n() % kAlignmentC)) {
+
+      return Status::kErrorMisalignedOperand;
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+    return 0;
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
+
+    // Determine grid shape
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(
+      args.problem_size,
+      {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+      args.batch_count);
+
+    // Initialize the Params structure
+    params_ = typename GemmKernel::Params{
+      args.problem_size,
+      grid_shape,
+      args.ref_A.non_const_ref(),
+      args.stride_A,
+      args.ref_B.non_const_ref(),
+      args.stride_B,
+      args.ref_C.non_const_ref(),
+      args.stride_C,
+      args.ref_D,
+      args.stride_D,
+      args.epilogue,
+      args.batch_count
+    };
+
+    return Status::kSuccess;
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+
+    params_.ref_A.reset(args.ref_A.non_const_ref().data());
+    params_.ref_B.reset(args.ref_B.non_const_ref().data());
+    params_.ref_C.reset(args.ref_C.non_const_ref().data());
+    params_.ref_D.reset(args.ref_D.data());
+
+    return Status::kSuccess;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+
+    ThreadblockSwizzle threadblock_swizzle;
+
+    dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
+    // XXX(Peter Han): prealod needs double warps in z direction
+    dim3 block(GemmKernel::kThreadCount, 1, kStages ? 1 : 2);
+
+    cudaError_t result;
+
+    int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
+    /// cudaFuncSetAttribute isn't supported under CUDA-8.0
+    // if (smem_size >= (48 << 10)) {
+    //   result = cudaFuncSetAttribute(Kernel<GemmKernel>,
+    //                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
+    //                                 smem_size);
+
+    //   if (result != cudaSuccess) {
+    //     return Status::kErrorInternal;
+    //   }
+
+    //   result = cudaFuncSetAttribute(
+    //       Kernel<GemmKernel>,
+    //       cudaFuncAttributePreferredSharedMemoryCarveout, 100);
+
+    //   if (result != cudaSuccess) {
+    //     return Status::kErrorInternal;
+    //   }
+    // }
+
+    cutlass::Kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params_);
+
+    result = cudaGetLastError();
+
+    return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(
+    Arguments const &args,
+    void *workspace = nullptr,
+    cudaStream_t stream = nullptr) {
+
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Parital specialization for column-major output exchanges problem size and operand.
+template <
+  /// Element type for A matrix operand
+  typename ElementA_,
+  /// Layout type for A matrix operand
+  typename LayoutA_,
+  /// Element type for B matrix operand
+  typename ElementB_,
+  /// Layout type for B matrix operand
+  typename LayoutB_,
+  /// Element type for C and D matrix operands
+  typename ElementC_,
+  /// Element type for internal accumulation
+  typename ElementAccumulator_,
+  /// Operator class tag
+  typename OperatorClass_,
+  /// Tag indicating architecture to tune for
+  typename ArchTag_,
+  /// Threadblock-level tile size (concept: GemmShape)
+  typename ThreadblockShape_,
+  /// Warp-level tile size (concept: GemmShape)
+  typename WarpShape_,
+  /// Warp-level tile size (concept: GemmShape)
+  typename InstructionShape_,
+  /// Epilogue output operator
+  typename EpilogueOutputOp_,
+  /// Threadblock-level swizzling operator
+  typename ThreadblockSwizzle_,
+  /// Number of stages used in the pipelined mainloop
+  int Stages,
+  /// Access granularity of A matrix in units of elements
+  int AlignmentA,
+  /// Access granularity of B matrix in units of elements
+  int AlignmentB,
+  typename Operator_
+>
+class GemmBatched<
+  ElementA_,
+  LayoutA_,
+  ElementB_,
+  LayoutB_,
+  ElementC_,
+  layout::ColumnMajor,
+  ElementAccumulator_,
+  OperatorClass_,
+  ArchTag_,
+  ThreadblockShape_,
+  WarpShape_,
+  InstructionShape_,
+  EpilogueOutputOp_,
+  ThreadblockSwizzle_,
+  Stages,
+  AlignmentA,
+  AlignmentB,
+  Operator_
+> {
+public:
+
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = layout::ColumnMajor;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static int const kStages = Stages;
+
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  static bool const kSplitKSerial = false;
+
+  //
+  using UnderlyingOperator = GemmBatched<
+    ElementB,
+    typename layout::LayoutTranspose<LayoutB>::type,
+    ElementA,
+    typename layout::LayoutTranspose<LayoutA>::type,
+    ElementC,
+    layout::RowMajor,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    kAlignmentB,
+    kAlignmentA
+  >;
+
+  using UnderlyingArguments = typename UnderlyingOperator::Arguments;
+  using GemmKernel = typename UnderlyingOperator::GemmKernel;
+
+  /// Argument structure
+  struct Arguments {
+
+    //
+    // Data members
+    //
+
+    GemmCoord problem_size;
+    TensorRef<ElementA const, LayoutA> ref_A;
+    int64_t stride_A;
+    TensorRef<ElementB const, LayoutB> ref_B;
+    int64_t stride_B;
+    TensorRef<ElementC const, LayoutC> ref_C;
+    int64_t stride_C;
+    TensorRef<ElementC, LayoutC> ref_D;
+    int64_t stride_D;
+    typename EpilogueOutputOp::Params epilogue;
+    int batch_count;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments() { }
+
+    /// Constructs an Arguments structure
+    CUTLASS_HOST_DEVICE
+    Arguments(
+      GemmCoord problem_size_,
+      TensorRef<ElementA const, LayoutA> ref_A_,
+      int64_t stride_A_,
+      TensorRef<ElementB const, LayoutB> ref_B_,
+      int64_t stride_B_,
+      TensorRef<ElementC const, LayoutC> ref_C_,
+      int64_t stride_C_,
+      TensorRef<ElementC, LayoutC> ref_D_,
+      int64_t stride_D_,
+      typename EpilogueOutputOp::Params epilogue_,
+      int batch_count_
+    ):
+      problem_size(problem_size_),
+      ref_A(ref_A_),
+      stride_A(stride_A_),
+      ref_B(ref_B_),
+      stride_B(stride_B_),
+      ref_C(ref_C_),
+      stride_C(stride_C_),
+      ref_D(ref_D_),
+      stride_D(stride_D_),
+      epilogue(epilogue_),
+      batch_count(batch_count_) { }
+  };
+
+private:
+
+  UnderlyingOperator underlying_operator_;
+
+public:
+
+  /// Constructs the GEMM.
+  GemmBatched() { }
+
+  /// Helper to construct a transposed equivalent for the underying GEMM operator
+  static UnderlyingArguments to_underlying_arguments(Arguments const &args) {
+    return UnderlyingArguments(
+      {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()},
+      {args.ref_B.data(), args.ref_B.stride(0)},
+      args.stride_B,
+      {args.ref_A.data(), args.ref_A.stride(0)},
+      args.stride_A,
+      {args.ref_C.data(), args.ref_C.stride(0)},
+      args.stride_C,
+      {args.ref_D.data(), args.ref_D.stride(0)},
+      args.stride_D,
+      args.epilogue,
+      args.batch_count
+    );
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+
+    return UnderlyingOperator::can_implement(to_underlying_arguments(args));
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+
+    return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.initialize(to_underlying_arguments(args), workspace);
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+
+    return underlying_operator_.update(to_underlying_arguments(args), workspace);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(
+    Arguments const &args,
+    void *workspace = nullptr,
+    cudaStream_t stream = nullptr) {
+
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace gemm
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

cat_files/gemm_device.h

@@ -0,0 +1,732 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/device_kernel.h"
+
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+#include "cutlass/gemm/kernel/gemm.h"
+
+#include "cutlass/gemm/kernel/default_gemm.h"
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace device {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/*! Gemm device-level operator. This is an interface to efficient CUTLASS GEMM kernels that may
+  be invoked from host code.
+
+  The contributions of this class are:
+
+    1. At compile time, it maps data types and high-level structural parameters onto
+       specific CUTLASS components.
+
+    2. At runtime, it maps logical arguments to GEMM problems to kernel parameters.
+
+    3. At runtime, it launches kernels on the device.
+
+  The intent is to provide a convenient mechanism for interacting with most plausible GEMM
+  configurations for each supported architecture. Consequently, not all parameters are exposed
+  to the top-level interface. Rather, sensible defaults at each level of the CUTLASS hierarchy
+  are selected to tradeoff simplicity of the interface with flexibility. We expect
+  most configurations to be specified at this level. Applications with more exotic requirements
+  may construct their kernels of interest using CUTLASS components at the threadblock, warp,
+  and thread levels of abstraction.
+
+  CUTLASS exposes computations using the functor design pattern in which objects compose some
+  internal state with an overloaded function call operator. This enables decoupling of
+  initialization from execution, possibly reducing overhead during steady state phases of
+  application execution.
+
+  CUTLASS device-level operators expose an Arguments structure encompassing each logical
+  input to the computation. This is distinct from the kernel-level Params structure pattern
+  which contains application-specific precomputed state needed by the device code.
+
+  Example of a CUTLASS GEMM operator implementing the functionality of cuBLAS's SGEMM NN
+  is as follows:
+
+    //
+    // Instantiate the CUTLASS GEMM operator.
+    //
+
+    cutlass::gemm::device::Gemm<
+      float,
+      cutlass::layout::ColumnMajor,
+      float,
+      cutlass::layout::ColumnMajor,
+      float,
+      cutlass::layout::ColumnMajor
+    > gemm_op;
+
+    //
+    // Launch the GEMM operation on the device
+    //
+
+    cutlass::Status status = gemm_op({
+      {m, n, k},                          // GemmCoord problem_size,
+      {A, lda},                           // TensorRef<float, layout::ColumnMajor> ref_A,
+      {B, ldb},                           // TensorRef<float, layout::ColumnMajor> ref_B,
+      {C, ldc},                           // TensorRef<float, layout::ColumnMajor> ref_C,
+      {D, ldd},                           // TensorRef<float, layout::ColumnMajor> ref_D,
+      {alpha, beta}                       // EpilogueOutputOp::Params epilogue_op_params
+    });
+
+
+  A simplified view of the template is listed below.
+
+    template <
+      /// Element type for A matrix operand
+      typename ElementA,
+
+      /// Layout type for A matrix operand
+      typename LayoutA,
+
+      /// Element type for B matrix operand
+      typename ElementB,
+
+      /// Layout type for B matrix operand
+      typename LayoutB,
+
+      /// Element type for C and D matrix operands
+      typename ElementC,
+
+      /// Layout type for C and D matrix operands
+      typename LayoutC,
+
+      /// Element type for internal accumulation
+      typename ElementAccumulator,
+
+      /// Operator class tag
+      typename OperatorClass,
+
+      /// Tag indicating architecture to tune for.  This is the minimum SM that
+      /// supports the intended feature. The device kernel can be built
+      /// targeting any SM larger than this number.
+      typename ArchTag,
+
+      /// Threadblock-level tile size (concept: GemmShape)
+      typename ThreadblockShape,
+
+      /// Warp-level tile size (concept: GemmShape)
+      typename WarpShape,
+
+      /// Warp-level tile size (concept: GemmShape)
+      typename InstructionShape,
+
+      /// Epilogue output operator
+      typename EpilogueOutputOp,
+
+      /// Threadblock-level swizzling operator
+      typename ThreadblockSwizzle,
+
+      /// Number of stages used in the pipelined mainloop
+      int Stages
+    >
+    class Gemm;
+*/
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_ = ElementC_,
+    /// Operator class tag
+    typename OperatorClass_ = arch::OpClassSimt,
+    /// Tag indicating architecture to tune for
+    typename ArchTag_ = arch::Sm61,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_ =
+        typename threadblock::GemmIdentityThreadblockSwizzle<>,
+    /// Number of stages used in the pipelined mainloop
+    int Stages =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kStages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentB,
+    /// If true, kernel supports split-K with serial reduction
+    bool SplitKSerial = false,
+    /// Operation performed by GEMM
+    typename Operator_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::Operator>
+class Gemm {
+ public:
+
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = LayoutC_;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  static bool const kSplitKSerial = SplitKSerial;
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
+
+  /// Define the kernel
+  using GemmKernel = typename kernel::DefaultGemm<
+    ElementA,
+    LayoutA,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    kAlignmentB,
+    ElementC,
+    LayoutC,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    kStages,
+    kSplitKSerial,
+    Operator
+  >::GemmKernel;
+
+  /// Argument structure
+  struct Arguments {
+
+    //
+    // Data members
+    //
+
+    GemmCoord problem_size;
+    TensorRef<ElementA const, LayoutA> ref_A;
+    TensorRef<ElementB const, LayoutB> ref_B;
+    TensorRef<ElementC const, LayoutC> ref_C;
+    TensorRef<ElementC, LayoutC> ref_D;
+    typename EpilogueOutputOp::Params epilogue;
+    int split_k_slices;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments(): problem_size(0, 0, 0), split_k_slices(1) {
+
+    }
+
+    /// Constructs an Arguments structure
+    CUTLASS_HOST_DEVICE
+    Arguments(
+      GemmCoord problem_size_,
+      TensorRef<ElementA const, LayoutA> ref_A_,
+      TensorRef<ElementB const, LayoutB> ref_B_,
+      TensorRef<ElementC const, LayoutC> ref_C_,
+      TensorRef<ElementC, LayoutC> ref_D_,
+      typename EpilogueOutputOp::Params epilogue_ =
+        typename EpilogueOutputOp::Params(),
+      int split_k_slices = 1
+    ):
+      problem_size(problem_size_),
+      ref_A(ref_A_),
+      ref_B(ref_B_),
+      ref_C(ref_C_),
+      ref_D(ref_D_),
+      epilogue(epilogue_),
+      split_k_slices(split_k_slices) {
+
+    }
+  };
+
+private:
+
+  /// Kernel parameters object
+  typename GemmKernel::Params params_;
+
+public:
+
+  /// Constructs the GEMM.
+  Gemm() { }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+
+    if (!kSplitKSerial && args.split_k_slices > 1) {
+      return Status::kErrorInvalidProblem;
+    }
+
+    Status status = GemmKernel::can_implement(
+      args.problem_size,
+      args.ref_A.non_const_ref(),
+      args.ref_B.non_const_ref(),
+      args.ref_C.non_const_ref(),
+      args.ref_D
+    );
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+
+    size_t bytes = 0;
+
+    // Determine grid shape
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape(
+      args.problem_size,
+      {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+      args.split_k_slices);
+
+    if (kSplitKSerial && args.split_k_slices > 1) {
+
+      bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n());
+    }
+
+    return bytes;
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
+
+    // Determine grid shape
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape(
+      args.problem_size,
+      {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
+      args.split_k_slices);
+
+    if (kSplitKSerial) {
+      if (args.split_k_slices > 1) {
+        if (!workspace) {
+          return Status::kErrorWorkspaceNull;
+        }
+
+        size_t bytes = get_workspace_size(args);
+
+        cudaError_t result = cudaMemsetAsync(workspace, 0, bytes, stream);
+
+        if (result != cudaSuccess) {
+          return Status::kErrorInternal;
+        }
+      }
+    }
+    else {
+
+      if (args.split_k_slices > 1) {
+        return Status::kErrorInvalidProblem;
+      }
+    }
+
+    // Initialize the Params structure
+    params_ = typename GemmKernel::Params{
+      args.problem_size,
+      grid_shape,
+      args.ref_A.non_const_ref(),
+      args.ref_B.non_const_ref(),
+      args.ref_C.non_const_ref(),
+      args.ref_D,
+      args.epilogue,
+      static_cast<int *>(workspace)
+    };
+
+    return Status::kSuccess;
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+
+    if (kSplitKSerial && args.split_k_slices > 1) {
+      if (!workspace) {
+        return Status::kErrorWorkspaceNull;
+      }
+    }
+
+    params_.ref_A.reset(args.ref_A.non_const_ref().data());
+    params_.ref_B.reset(args.ref_B.non_const_ref().data());
+    params_.ref_C.reset(args.ref_C.non_const_ref().data());
+    params_.ref_D.reset(args.ref_D.data());
+    params_.output_op = args.epilogue;
+    params_.semaphore = static_cast<int *>(workspace);
+
+    return Status::kSuccess;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+
+    ThreadblockSwizzle threadblock_swizzle;
+
+    dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape);
+    // XXX(Peter Han): prealod needs double warps in z direction
+    dim3 block(GemmKernel::kThreadCount, 1, kStages ? 1 : 2);
+
+    cudaError_t result;
+
+    int smem_size = int(sizeof(typename GemmKernel::SharedStorage));
+    /// cudaFuncSetAttribute isn't supported under CUDA-8.0
+    // if (smem_size >= (48 << 10)) {
+    //   result = cudaFuncSetAttribute(Kernel<GemmKernel>,
+    //                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
+    //                                 smem_size);
+
+    //   if (result != cudaSuccess) {
+    //     return Status::kErrorInternal;
+    //   }
+
+    //   result = cudaFuncSetAttribute(
+    //       Kernel<GemmKernel>,
+    //       cudaFuncAttributePreferredSharedMemoryCarveout, 100);
+
+    //   if (result != cudaSuccess) {
+    //     return Status::kErrorInternal;
+    //   }
+    // }
+
+    cutlass::Kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params_);
+
+    result = cudaGetLastError();
+
+    return result == cudaSuccess ? Status::kSuccess : Status::kErrorInternal;
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(
+    Arguments const &args,
+    void *workspace = nullptr,
+    cudaStream_t stream = nullptr) {
+
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Parital specialization for column-major output exchanges problem size and operand.
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_,
+    /// Operator class tag
+    typename OperatorClass_,
+    /// Tag indicating architecture to tune for
+    typename ArchTag_,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB,
+    /// If true, kernel supports split-K as a serial reduction
+    bool SplitKSerial,
+    /// Operation performed by GEMM
+    typename Operator_>
+class Gemm<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
+           layout::ColumnMajor,  // partially specialized on LayoutC
+           ElementAccumulator_, OperatorClass_, ArchTag_, ThreadblockShape_,
+           WarpShape_, InstructionShape_, EpilogueOutputOp_,
+           ThreadblockSwizzle_, Stages, AlignmentA, AlignmentB, SplitKSerial,
+           Operator_> {
+ public:
+
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = layout::ColumnMajor;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
+  static bool const kSplitKSerial = SplitKSerial;
+
+  using UnderlyingOperator = Gemm<
+    ElementB,
+    typename layout::LayoutTranspose<LayoutB>::type,
+    ElementA,
+    typename layout::LayoutTranspose<LayoutA>::type,
+    ElementC,
+    layout::RowMajor,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    kAlignmentB,
+    kAlignmentA,
+    SplitKSerial,
+    Operator
+  >;
+
+  using UnderlyingArguments = typename UnderlyingOperator::Arguments;
+  using GemmKernel = typename UnderlyingOperator::GemmKernel;
+  static int const kAlignmentC = UnderlyingOperator::kAlignmentC;
+
+  /// Argument structure
+  struct Arguments {
+
+    //
+    // Data members
+    //
+
+    GemmCoord problem_size;
+    TensorRef<ElementA const, LayoutA> ref_A;
+    TensorRef<ElementB const, LayoutB> ref_B;
+    TensorRef<ElementC const, LayoutC> ref_C;
+    TensorRef<ElementC, LayoutC> ref_D;
+    typename EpilogueOutputOp::Params epilogue;
+    int split_k_slices;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments() { }
+
+    /// Constructs an Arguments structure
+    CUTLASS_HOST_DEVICE
+    Arguments(
+      GemmCoord problem_size_,
+      TensorRef<ElementA const, LayoutA> ref_A_,
+      TensorRef<ElementB const, LayoutB> ref_B_,
+      TensorRef<ElementC const, LayoutC> ref_C_,
+      TensorRef<ElementC, LayoutC> ref_D_,
+      typename EpilogueOutputOp::Params epilogue_ =
+        typename EpilogueOutputOp::Params(),
+      int split_k_slices = 1
+    ):
+      problem_size(problem_size_),
+      ref_A(ref_A_),
+      ref_B(ref_B_),
+      ref_C(ref_C_),
+      ref_D(ref_D_),
+      epilogue(epilogue_),
+      split_k_slices(split_k_slices) { }
+  };
+
+private:
+
+  UnderlyingOperator underlying_operator_;
+
+public:
+
+  /// Constructs the GEMM.
+  Gemm() { }
+
+  /// Helper to construct a transposed equivalent for the underying GEMM operator
+  static UnderlyingArguments to_underlying_arguments(Arguments const &args) {
+    return UnderlyingArguments(
+      {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()},
+      {args.ref_B.data(), args.ref_B.stride(0)},
+      {args.ref_A.data(), args.ref_A.stride(0)},
+      {args.ref_C.data(), args.ref_C.stride(0)},
+      {args.ref_D.data(), args.ref_D.stride(0)},
+      args.epilogue,
+      args.split_k_slices
+    );
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+
+    return UnderlyingOperator::can_implement(to_underlying_arguments(args));
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+
+    return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream);
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+
+    return underlying_operator_.update(to_underlying_arguments(args), workspace);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(
+    Arguments const &args,
+    void *workspace = nullptr,
+    cudaStream_t stream = nullptr) {
+
+    Status status = initialize(args, workspace, stream);
+
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace gemm
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

cat_files/gemm_universal.h

@@ -0,0 +1,376 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/arch/arch.h"
+#include "cutlass/device_kernel.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+#include "cutlass/gemm/kernel/gemm_universal.h"
+
+#include "cutlass/gemm/kernel/default_gemm_universal.h"
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+#include "cutlass/gemm/device/gemm_universal_base.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace device {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/*! 
+  The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and 
+  batched array variants.
+*/
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_ = ElementC_,
+    /// Operator class tag
+    typename OperatorClass_ = arch::OpClassSimt,
+    /// Tag indicating architecture to tune for.  This is the minimum SM that
+    /// supports the intended feature. The device kernel can be built
+    /// targeting any SM larger than this number.
+    typename ArchTag_ = arch::Sm61,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
+    /// Number of stages used in the pipelined mainloop
+    int Stages =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kStages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB =
+        DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
+                                 ElementC_, ElementAccumulator_>::kAlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_ = typename DefaultGemmConfiguration<
+        OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator_>::Operator,
+    /// Complex elementwise transformation on A operand
+    ComplexTransform TransformA = ComplexTransform::kNone,
+    /// Complex elementwise transformation on B operand
+    ComplexTransform TransformB = ComplexTransform::kNone
+>
+class GemmUniversal : 
+  GemmUniversalBase<
+    typename kernel::DefaultGemmUniversal<
+      ElementA_,
+      LayoutA_,
+      TransformA,
+      AlignmentA,
+      ElementB_,
+      LayoutB_,
+      TransformB,
+      AlignmentB,
+      ElementC_,
+      LayoutC_,
+      ElementAccumulator_,
+      OperatorClass_,
+      ArchTag_,
+      ThreadblockShape_,
+      WarpShape_,
+      InstructionShape_,
+      EpilogueOutputOp_,
+      ThreadblockSwizzle_,
+      Stages,
+      Operator_
+    >::GemmKernel
+  > {
+
+ public:
+
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  static ComplexTransform const kTransformA = TransformA;
+  static ComplexTransform const kTransformB = TransformB;
+
+  using Base = GemmUniversalBase<
+    typename kernel::DefaultGemmUniversal<
+      ElementA_,
+      LayoutA_,
+      TransformA,
+      AlignmentA,
+      ElementB_,
+      LayoutB_,
+      TransformB,
+      AlignmentB,
+      ElementC_,
+      LayoutC_,
+      ElementAccumulator_,
+      OperatorClass_,
+      ArchTag_,
+      ThreadblockShape_,
+      WarpShape_,
+      InstructionShape_,
+      EpilogueOutputOp_,
+      ThreadblockSwizzle_,
+      Stages,
+      Operator_
+    >::GemmKernel
+  >;
+
+  using Arguments = typename Base::Arguments;
+  using GemmKernel = typename Base::GemmKernel;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Parital specialization for column-major output exchanges problem size and operand.
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_,
+    /// Operator class tag
+    typename OperatorClass_,
+    /// Tag indicating architecture to tune for.  This is the minimum SM that
+    /// supports the intended feature. The device kernel can be built
+    /// targeting any SM larger than this number.
+    typename ArchTag_,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_,
+    /// Complex elementwise transformation on A operand
+    ComplexTransform TransformA,
+    /// Complex elementwise transformation on B operand
+    ComplexTransform TransformB>
+class GemmUniversal<ElementA_, LayoutA_, ElementB_, LayoutB_, ElementC_,
+           layout::ColumnMajor,  // partially specialized on LayoutC
+           ElementAccumulator_, OperatorClass_, ArchTag_, ThreadblockShape_,
+           WarpShape_, InstructionShape_, EpilogueOutputOp_,
+           ThreadblockSwizzle_, Stages, AlignmentA, AlignmentB,
+           Operator_, TransformA, TransformB> {
+ public:
+
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = layout::ColumnMajor;
+  using TensorRefC = TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static ComplexTransform const kTransformA = TransformA;
+  static ComplexTransform const kTransformB = TransformB;
+
+  using UnderlyingOperator = typename GemmUniversal< 
+    ElementB,
+    typename layout::LayoutTranspose<LayoutB>::type,
+    ElementA,
+    typename layout::LayoutTranspose<LayoutA>::type,
+    ElementC,
+    layout::RowMajor,    
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    kAlignmentB,
+    kAlignmentA,
+    Operator,
+    kTransformB,
+    kTransformA
+  >::Base;
+
+  using GemmKernel = typename UnderlyingOperator::GemmKernel;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+
+  /// Argument structure
+  using Arguments = typename UnderlyingOperator::Arguments;
+
+private:
+
+  UnderlyingOperator underlying_operator_;
+
+public:
+
+  /// Constructs the GEMM.
+  GemmUniversal() { }
+
+  /// Helper to construct a transposed equivalent for the underying GEMM operator
+  static Arguments to_underlying_arguments(Arguments const &args) {
+    return args.transposed_problem();
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status can_implement(Arguments const &args) {
+
+    return UnderlyingOperator::can_implement(to_underlying_arguments(args));
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+    
+    return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args));
+  }
+
+  /// Computes the grid shape
+  static dim3 get_grid_shape(Arguments const &args) { 
+    return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args));
+  }
+
+  /// Computes the maximum number of active blocks per multiprocessor
+  static int maximum_active_blocks(int smem_capacity = -1) {
+    return UnderlyingOperator::maximum_active_blocks(smem_capacity);
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status initialize(Arguments const &args, void *workspace = nullptr, cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream);
+  }
+
+  /// Lightweight update given a subset of arguments
+  Status update(Arguments const &args, void *workspace = nullptr) {
+
+    return underlying_operator_.update(to_underlying_arguments(args), workspace);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status run(cudaStream_t stream = nullptr) {
+
+    return underlying_operator_.run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(cudaStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  Status operator()(
+    Arguments const &args, 
+    void *workspace = nullptr, 
+    cudaStream_t stream = nullptr) {
+    
+    Status status = initialize(args, workspace, stream);
+    
+    if (status == Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace gemm
+} // namespace cutlass
+
+////////////////////////////////////////////////////////////////////////////////

cat_files/mma_cu10.h

@@ -0,0 +1,394 @@
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Matrix Multiply for BigIsland 1st generation
+*/
+
+#pragma once
+
+#include "cutlass/arch/mma.h"
+
+#include "cutlass/layout/matrix.h"
+#include "cutlass/gemm/gemm.h"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace arch {
+
+/// BigIsland Tensor Core tile format - EM orinted vector type definitions
+/// fp32
+typedef float v4float_t __attribute__((ext_vector_type(4)));
+/// s32
+typedef int32_t v4int32_t __attribute__((ext_vector_type(4)));
+/// u32
+typedef uint32_t v4uint32_t __attribute__((ext_vector_type(4)));
+/// fp16
+typedef uint16_t v4half_t __attribute__((ext_vector_type(4)));
+/// bf16
+typedef uint16_t v4bfloat16_t __attribute__((ext_vector_type(4)));
+/// s8
+typedef int8_t v4int8_t __attribute__((ext_vector_type(4)));
+/// u8
+typedef uint8_t v4uint8_t __attribute__((ext_vector_type(4)));
+
+////////////////////////////////////////////////////////////////////////////////
+//
+// Matrix multiply accumulate 161616 - U32 accumulation
+//
+////////////////////////////////////////////////////////////////////////////////
+
+/// Matrix multiply-add operation: U32 = U8 * U8 + U32
+template <typename LayoutA, typename LayoutB, typename LayoutC>
+struct Mma<
+  gemm::GemmShape<16, 16, 16>,
+  64,
+  uint8_t,
+  LayoutA,
+  uint8_t,
+  LayoutB,
+  uint32_t,
+  LayoutC,
+  OpMultiplyAdd> {
+
+  using Shape = gemm::GemmShape<16, 16, 16>;
+
+  using ElementA = uint8_t;
+  using FragmentA = Array<uint8_t, 4>;
+
+  using ElementB = uint8_t;
+  using FragmentB = Array<uint8_t, 4>;
+
+  using ElementC = uint;
+  using FragmentC = Array<uint, 4>;
+
+  using Operator = OpMultiplyAdd;
+  using ArchTag = arch::Cu10;
+
+  CUTLASS_HOST_DEVICE
+  void operator()(
+    FragmentC &d,
+    FragmentA const &a,
+    FragmentB const &b,
+    FragmentC const &c
+  ) const {
+#if CUTLASS_ARCH_CU10_SUPPORTED
+    v4uint8_t src_A;
+    v4uint8_t src_B;
+    v4uint32_t src_C;
+    v4uint32_t dst_D;
+
+    src_A[0] = a[0];
+    src_A[1] = a[1];
+    src_A[2] = a[2];
+    src_A[3] = a[3];
+    src_B[0] = b[0];
+    src_B[1] = b[1];
+    src_B[2] = b[2];
+    src_B[3] = b[3];
+    src_C[0] = c[0];
+    src_C[1] = c[1];
+    src_C[2] = c[2];
+    src_C[3] = c[3];
+
+    dst_D = __ivcorex_matrix_mad_u32x4_u8x4(src_A, src_B, src_C);
+
+    d[0] = dst_D[0];
+    d[1] = dst_D[1];
+    d[2] = dst_D[2];
+    d[3] = dst_D[3];
+#else
+    assert(0);
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+//
+// Matrix multiply accumulate 161616 - S32 accumulation
+//
+////////////////////////////////////////////////////////////////////////////////
+
+/// Matrix multiply-add operation: S32 = S8 * S8 + S32
+template <typename LayoutA, typename LayoutB, typename LayoutC>
+struct Mma<
+  gemm::GemmShape<16, 16, 16>,
+  64,
+  int8_t,
+  LayoutA,
+  int8_t,
+  LayoutB,
+  int,
+  LayoutC,
+  OpMultiplyAdd> {
+
+  using Shape = gemm::GemmShape<16, 16, 16>;
+
+  using ElementA = int8_t;
+  using FragmentA = Array<int8_t, 4>;
+
+  using ElementB = int8_t;
+  using FragmentB = Array<int8_t, 4>;
+
+  using ElementC = int;
+  using FragmentC = Array<int, 4>;
+
+  using Operator = OpMultiplyAdd;
+  using ArchTag = arch::Cu10;
+
+  CUTLASS_HOST_DEVICE
+  void operator()(
+    FragmentC &d,
+    FragmentA const &a,
+    FragmentB const &b,
+    FragmentC const &c
+  ) const {
+#if CUTLASS_ARCH_CU10_SUPPORTED
+    v4int8_t src_A;
+    v4int8_t src_B;
+    v4int32_t src_C;
+    v4int32_t dst_D;
+
+    src_A[0] = a[0];
+    src_A[1] = a[1];
+    src_A[2] = a[2];
+    src_A[3] = a[3];
+    src_B[0] = b[0];
+    src_B[1] = b[1];
+    src_B[2] = b[2];
+    src_B[3] = b[3];
+    src_C[0] = c[0];
+    src_C[1] = c[1];
+    src_C[2] = c[2];
+    src_C[3] = c[3];
+
+    dst_D = __ivcorex_matrix_mad_i32x4_i8x4(src_A, src_B, src_C);
+
+    d[0] = dst_D[0];
+    d[1] = dst_D[1];
+    d[2] = dst_D[2];
+    d[3] = dst_D[3];
+#else
+    assert(0);
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+//
+// Matrix multiply accumulate 161616 - FP32 accumulation
+//
+////////////////////////////////////////////////////////////////////////////////
+
+/// Matrix multiply-add operation: FP32 = FP16 * FP16 + FP32
+template <typename LayoutA, typename LayoutB, typename LayoutC>
+struct Mma<
+  gemm::GemmShape<16, 16, 16>,
+  64,
+  cutlass::half_t,
+  LayoutA,
+  cutlass::half_t,
+  LayoutB,
+  float,
+  LayoutC,
+  OpMultiplyAdd> {
+
+  using Shape = gemm::GemmShape<16, 16, 16>;
+
+  using ElementA = cutlass::half_t;
+  using FragmentA = Array<half_t, 4>;
+
+  using ElementB = cutlass::half_t;
+  using FragmentB = Array<half_t, 4>;
+
+  using ElementC = float;
+  using FragmentC = Array<float, 4>;
+
+  using Operator = OpMultiplyAdd;
+  using ArchTag = arch::Cu10;
+
+  CUTLASS_HOST_DEVICE
+  void operator()(
+    FragmentC &d,
+    FragmentA const &a,
+    FragmentB const &b,
+    FragmentC const &c
+  ) const {
+    v4half_t src_A;
+    v4half_t src_B;
+    v4float_t src_C;
+    v4float_t dst_D;
+
+    src_A[0] = half_t(a[0]).storage;
+    src_A[1] = half_t(a[1]).storage;
+    src_A[2] = half_t(a[2]).storage;
+    src_A[3] = half_t(a[3]).storage;
+    src_B[0] = half_t(b[0]).storage;
+    src_B[1] = half_t(b[1]).storage;
+    src_B[2] = half_t(b[2]).storage;
+    src_B[3] = half_t(b[3]).storage;
+    src_C[0] = c[0];
+    src_C[1] = c[1];
+    src_C[2] = c[2];
+    src_C[3] = c[3];
+
+    dst_D = __ivcorex_matrix_mad_f32x4_f16x4(src_A, src_B, src_C);
+#if 0
+if(threadIdx.x == 0)
+printf(
+  ">>> After\n"
+  "A: %f, %f, %f, %f\n"
+  "B: %f, %f, %f, %f\n"
+  "C: %f, %f, %f, %f\n"
+  "D: %f, %f, %f, %f\n\n",
+  float(a[0]), float(a[1]), float(a[2]), float(a[3]),
+  float(b[0]), float(b[1]), float(b[2]), float(b[3]),
+  float(src_C[0]), float(src_C[1]), float(src_C[2]), float(src_C[3]),
+  float(d[0]), float(d[1]), float(d[2]), float(d[3])
+);
+#endif
+
+    d[0] = dst_D[0];
+    d[1] = dst_D[1];
+    d[2] = dst_D[2];
+    d[3] = dst_D[3];
+
+  }
+};
+
+/// Matrix multiply-add operation: FP32 = BF16 * BF16 + FP32
+template <typename LayoutA, typename LayoutB, typename LayoutC>
+struct Mma<
+  gemm::GemmShape<16, 16, 16>,
+  64,
+  bfloat16_t,
+  LayoutA,
+  bfloat16_t,
+  LayoutB,
+  float,
+  LayoutC,
+  OpMultiplyAdd> {
+
+  using Shape = gemm::GemmShape<16, 16, 16>;
+
+  using ElementA = bfloat16_t;
+  using FragmentA = Array<bfloat16_t, 4>;
+
+  using ElementB = bfloat16_t;
+  using FragmentB = Array<bfloat16_t, 4>;
+
+  using ElementC = float;
+  using FragmentC = Array<float, 4>;
+
+  using Operator = OpMultiplyAdd;
+  using ArchTag = arch::Cu10;
+
+  CUTLASS_HOST_DEVICE
+  void operator()(
+    FragmentC &d,
+    FragmentA const &a,
+    FragmentB const &b,
+    FragmentC const &c
+  ) const {
+    v4bfloat16_t src_A;
+    v4bfloat16_t src_B;
+    v4float_t src_C;
+    v4float_t dst_D;
+
+    src_A[0] = bfloat16_t(a[0]).storage;
+    src_A[1] = bfloat16_t(a[1]).storage;
+    src_A[2] = bfloat16_t(a[2]).storage;
+    src_A[3] = bfloat16_t(a[3]).storage;
+    src_B[0] = bfloat16_t(b[0]).storage;
+    src_B[1] = bfloat16_t(b[1]).storage;
+    src_B[2] = bfloat16_t(b[2]).storage;
+    src_B[3] = bfloat16_t(b[3]).storage;
+    src_C[0] = c[0];
+    src_C[1] = c[1];
+    src_C[2] = c[2];
+    src_C[3] = c[3];
+#if __clang_major__ >= 16
+    dst_D = __ivcorex_matrix_mad_f32x4_bf16x4(src_A, src_B, src_C);
+#else
+    dst_D = __ivcorex_matrix_mad_f32_bf16(src_A, src_B, src_C);
+#endif
+    d[0] = dst_D[0];
+    d[1] = dst_D[1];
+    d[2] = dst_D[2];
+    d[3] = dst_D[3];
+  }
+};
+
+/// Matrix multiply-add operation: FP32 = FP32 * FP32 + FP32
+template <typename LayoutA, typename LayoutB, typename LayoutC>
+struct Mma<
+  gemm::GemmShape<16,16,16>,
+  64,
+  float,
+  LayoutA,
+  float,
+  LayoutB,
+  float,
+  LayoutC,
+  OpMultiplyAdd> {
+
+  using Shape = gemm::GemmShape<16,16,16>;
+
+  using ElementA = float;
+  using FragmentA = Array<float, 4>;
+
+  using ElementB = float;
+  using FragmentB = Array<float, 4>;
+
+  using ElementC = float;
+  using FragmentC = Array<float, 4>;
+
+  using Operator = OpMultiplyAdd;
+  using ArchTag = arch::Cu10;
+
+  CUTLASS_HOST_DEVICE
+  void operator()(
+    FragmentC &d,
+    FragmentA const &a,
+    FragmentB const &b,
+    FragmentC const &c
+  ) const {
+    v4float_t src_A;
+    v4float_t src_B;
+    v4float_t src_C;
+    v4float_t dst_D;
+
+    src_A[0] = a[0];
+    src_A[1] = a[1];
+    src_A[2] = a[2];
+    src_A[3] = a[3];
+    src_B[0] = b[0];
+    src_B[1] = b[1];
+    src_B[2] = b[2];
+    src_B[3] = b[3];
+    src_C[0] = c[0];
+    src_C[1] = c[1];
+    src_C[2] = c[2];
+    src_C[3] = c[3];
+
+    dst_D = __ivcorex_matrix_mad_f32x4_f32x4(src_A, src_B, src_C);
+
+    d[0] = dst_D[0];
+    d[1] = dst_D[1];
+    d[2] = dst_D[2];
+    d[3] = dst_D[3];
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+}
+}

cat_files/mma_tensor_op.h

@@ -0,0 +1,382 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Templates implementing warp-level matrix multiply-accumulate operations targeting
+      Tensor Cores.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/array.h"
+#include "cutlass/platform/platform.h"
+
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/matrix_shape.h"
+
+#include "cutlass/arch/mma.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/warp/mma.h"
+
+#include "cutlass/gemm/warp/mma_tensor_op_policy.h"
+#include "cutlass/gemm/warp/mma_tensor_op_tile_iterator.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <typename T, typename S, int N, FloatRoundStyle Round>
+struct ConvertAndPack {
+
+  using Converter = NumericArrayConverter<T, S, N, Round>;
+
+  CUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<S, N> const &source) {
+    Converter converter;
+
+    return converter(source);
+  }
+};
+
+template <typename T, int N, FloatRoundStyle Round>
+struct ConvertAndPack<T, T, N, Round> {
+
+  CUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &source) {
+		return source;
+  }
+};
+
+template <int N, FloatRoundStyle Round>
+struct ConvertAndPack<bfloat16_t, float, N, Round> {
+
+  using Converter = NumericArrayConverter<bfloat16_t, float, N, Round>;
+
+  CUTLASS_HOST_DEVICE
+  Array<bfloat16_t, N> operator()(Array<float, N> const &source) {
+    Converter converter;
+
+    Array<float, N> tmp;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      int idx = (((i << 1) & 2) | ((i >> 1) & 1) | (i & 0xfffffffc));
+      tmp[i] = source[idx];
+    }
+
+    return converter(tmp);
+  }
+};
+
+template <int N, FloatRoundStyle Round>
+struct ConvertAndPack<half_t, float, N, Round> {
+
+  using Converter = NumericArrayConverter<half_t, float, N, Round>;
+
+  CUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<float, N> const &source) {
+    Converter converter;
+
+    Array<float, N> tmp;
+
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      int idx = (((i << 1) & 2) | ((i >> 1) & 1) | (i & 0xfffffffc));
+      tmp[i] = source[idx];
+    }
+
+    return converter(tmp);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions.
+template <
+  /// Size of the Gemm problem - concept: gemm::GemmShape<>
+  typename Shape_,
+  /// Data type of A elements
+  typename ElementA_,
+  /// Layout of A matrix (concept: MatrixLayout)
+  typename LayoutA_,
+  /// Data type of B elements
+  typename ElementB_,
+  /// Layout of B matrix (concept: MatrixLayout)
+  typename LayoutB_,
+  /// Element type of C matrix
+  typename ElementC_,
+  /// Layout of C matrix (concept: MatrixLayout)
+  typename LayoutC_,
+  /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy)
+  typename Policy_,
+  /// Number of partitions along K dimension
+  int PartitionsK_ = 1,
+  /// Store the accumulators in row major or column major.
+  /// Iluvatar Tensor Core always stores accumulators in row major
+  bool AccumulatorsInRowMajor = true,
+  /// Used for partial specialization
+  typename Enable = bool
+>
+class MmaTensorOp {
+public:
+  /// Shape of warp-level matrix operation (concept: GemmShape)
+  using Shape = Shape_;
+
+  /// Data type of multiplicand A
+  using ElementA = ElementA_;
+
+  /// Layout of multiplicand A
+  using LayoutA = LayoutA_;
+
+  /// Data type of multiplicand B
+  using ElementB = ElementB_;
+
+  /// Layout of multiplicand B
+  using LayoutB = LayoutB_;
+
+  /// Data type of accumulator matrix C
+  using ElementC = ElementC_;
+
+  /// Layout of accumulator matrix C
+  using LayoutC = LayoutC_;
+
+  /// Shape of the warp in units of thread (concept: MmaLanePolicySimt)
+  using Policy = Policy_;
+
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  using ArchMmaOperator = typename Policy::Operator;
+
+  /// Architecture tag from underlying instruction
+  using ArchTag = typename ArchMmaOperator::ArchTag;
+
+  /// Indicates class of matrix operator
+  using OperatorClass = arch::OpClassTensorOp;
+
+  /// Shape of underlying instruction
+  using InstructionShape = typename ArchMmaOperator::Shape;
+
+  /// Complex transform on A operand
+  static ComplexTransform const kTransformA = ComplexTransform::kNone;
+
+  /// Complex transform on B operand
+  static ComplexTransform const kTransformB = ComplexTransform::kNone;
+
+  /// Number of threads participating in warp-level matrix product
+  static int const kThreadCount = NUM_THREADS_PER_WARP;
+
+  /// Number of partitions along K dimension
+  static int const kPartitionsK = PartitionsK_;
+
+public:
+  /// FIXME(Peter Han): workaround to adapt to simt epilogue, need to remove
+  struct ThreadMma {
+    using ElementC = ElementC;
+  };
+
+  /// Iterates over the A operand in memory
+  using IteratorA = MmaTensorOpMultiplicandTileIterator<
+    MatrixShape<Shape::kM, Policy::Operator::Shape::kK>,
+    Operand::kA,
+    ElementA,
+    LayoutA,
+    InstructionShape,
+    kThreadCount,
+    kPartitionsK>;
+
+  /// Storage for A tile
+  using FragmentA = typename IteratorA::Fragment;
+
+  /// Storage for transformed A tile
+  using TransformedFragmentA =
+      Array<typename ArchMmaOperator::ElementA, FragmentA::kElements>;
+
+  /// Iterates over the B operand in memory
+  using IteratorB = MmaTensorOpMultiplicandTileIterator<
+    MatrixShape<Policy::Operator::Shape::kK, Shape::kN>,
+    Operand::kB,
+    ElementB,
+    LayoutB,
+    InstructionShape,
+    kThreadCount,
+    kPartitionsK>;
+
+  /// Storage for B tile
+  using FragmentB = typename IteratorB::Fragment;
+
+  /// Storage for transformed B tile
+  using TransformedFragmentB =
+      Array<typename ArchMmaOperator::ElementB, FragmentB::kElements>;
+
+  /// Iterates over the C operand in memory
+  using IteratorC = MmaTensorOpAccumulatorTileIterator<
+    MatrixShape<Shape::kM, Shape::kN>,
+    ElementC,
+    LayoutC,
+    InstructionShape>;
+
+  /// Storage for C tile
+  using FragmentC = typename IteratorC::Fragment;
+
+  static_assert(
+    !(Shape::kM % Policy::Operator::Shape::kM) &&
+    !(Shape::kN % Policy::Operator::Shape::kN) &&
+    !(Shape::kK % Policy::Operator::Shape::kK),
+    "Shape of warp-level Mma must be divisible by operator shape.");
+
+  using MmaIterations = gemm::GemmShape<
+    (Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM,
+    (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN,
+    InstructionShape::kK / Policy::Operator::Shape::kK
+  >;
+
+public:
+
+  /// Underlying matrix multiply operator (concept: arch::Mma)
+  ArchMmaOperator mma;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Ctor
+  CUTLASS_DEVICE
+  MmaTensorOp() {}
+
+  /// Performs a warp-level matrix multiply-accumulate operation
+  CUTLASS_DEVICE
+  void operator()(
+    FragmentC &D,
+    TransformedFragmentA const &A,
+    TransformedFragmentB const &B,
+    FragmentC const &C
+  ) const {
+
+    using MmaOperandA = typename ArchMmaOperator::FragmentA;
+    using MmaOperandB = typename ArchMmaOperator::FragmentB;
+    using MmaOperandC = typename ArchMmaOperator::FragmentC;
+
+    D = C;
+
+    MmaOperandA const *ptr_A = reinterpret_cast<MmaOperandA const *>(&A);
+    MmaOperandB const *ptr_B = reinterpret_cast<MmaOperandB const *>(&B);
+    MmaOperandC *ptr_D = reinterpret_cast<MmaOperandC *>(&D);
+
+    // Serpentine visitation order maximizing reuse of Rb
+    CUTLASS_PRAGMA_UNROLL
+    for (int k = 0; k < MmaIterations::kK; ++k) {
+      CUTLASS_PRAGMA_UNROLL
+      for (int m = 0; m < MmaIterations::kM; ++m) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int n = 0; n < MmaIterations::kN; ++n) {
+          int n_serpentine = ((m % 2) ? (MmaIterations::kN - 1 - n) : n);
+
+          /// assume A is column-major in VRF, B is row-major in VRF
+          if(AccumulatorsInRowMajor) {
+            mma(
+              ptr_D[n_serpentine + m * MmaIterations::kN],
+              ptr_A[m + k * MmaIterations::kM],
+              ptr_B[n_serpentine + k * MmaIterations::kN],
+              ptr_D[n_serpentine + m * MmaIterations::kN]);
+          } else {
+            mma(
+              ptr_D[m + n_serpentine * MmaIterations::kM],
+              ptr_A[m + k * MmaIterations::kM],
+              ptr_B[n_serpentine + k * MmaIterations::kN],
+              ptr_D[m + n_serpentine * MmaIterations::kM]);
+          }
+        }
+      }
+    }
+  }
+
+  /// Transform the mma operands to the required types
+  CUTLASS_DEVICE
+  void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B,
+                 FragmentA const &A, FragmentB const &B) const {
+
+    //
+    // Define conversions from source type to instruction type
+    //
+    FloatRoundStyle const kRoundA =
+        PreferredRoundingMode<typename ArchMmaOperator::ElementA,
+                              ElementA>::kRound;
+    FloatRoundStyle const kRoundB =
+        PreferredRoundingMode<typename ArchMmaOperator::ElementB,
+                              ElementB>::kRound;
+    detail::ConvertAndPack<typename ArchMmaOperator::ElementA, ElementA,
+                            FragmentA::kElements / 2, kRoundA>
+          convert_A;
+    NumericArrayConverter<typename ArchMmaOperator::ElementB, ElementB,
+                            FragmentB::kElements, kRoundB>
+          convert_B;
+    Array<ElementA, FragmentA::kElements / 2> const *ptr_A =
+          reinterpret_cast<Array<ElementA, FragmentA::kElements / 2> const *>(&A);
+    Array<typename ArchMmaOperator::ElementA, FragmentA::kElements / 2> *
+          ptr_dst_A = reinterpret_cast<Array<typename ArchMmaOperator::ElementA,
+                                             FragmentA::kElements / 2> *>(&dst_A);
+
+    dst_B = convert_B(B);
+
+    ptr_dst_A[0] = convert_A(ptr_A[0]);
+    ptr_dst_A[1] = convert_A(ptr_A[1]);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace gemm
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

cat_files/mma_tensor_op_policy.h

@@ -0,0 +1,71 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/***************************************************************************************************
+* Copyright (c) 2021 Iluvatar CoreX. All rights reserved.
+* Copyright Declaration: This software, including all of its code and documentation,
+* except for the third-party software it contains, is a copyrighted work of Shanghai Iluvatar CoreX
+* Semiconductor Co., Ltd. and its affiliates ("Iluvatar CoreX") in accordance with the PRC Copyright
+* Law and relevant international treaties, and all rights contained therein are enjoyed by Iluvatar
+* CoreX. No user of this software shall have any right, ownership or interest in this software and
+* any use of this software shall be in compliance with the terms and conditions of the End User
+* License Agreement.
+ **************************************************************************************************/
+
+/*! \file
+    \brief Policy describing implementation details of warp-level GEMM targeting Tensor Cores.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/matrix_shape.h"
+#include "cutlass/gemm/gemm.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace warp {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Policy
+template <
+  typename Operator_,        ///< hardware instruction(s) performing TensorOp (concept: arch::Mma)
+  typename OpDelta_          ///< distance between operations (concept: MatrixShape)
+>
+struct MmaTensorOpPolicy {
+
+  using Operator = Operator_;    ///< hardware instruction(s) performing TensorOp (concept: arch::Mma)
+  using OpDelta = OpDelta_;      ///< distance between operations (concept: MatrixShape)
+  using MmaShape = typename Operator::Shape;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace warp
+} // namespace gemm
+} // namespace cutlass

cat_files/symbol_dumps/sym_ixpkg__C.cpython-310-x86_64-linux-gnu.so.txt

@@ -0,0 +1,6 @@
+000000000008ed90 T PyInit__C
+000000000009af00 T _ZSt15get_new_handlerv
+000000000009ad80 T _ZdlPvSt11align_val_t
+000000000009ad90 T _ZnwmSt11align_val_t
+000000000009af70 T _fini
+0000000000019000 T _init

cat_files/symbol_dumps/sym_ixpkg__ixformer_torch.cpython-310-x86_64-linux-gnu.so.txt

@@ -0,0 +1,49 @@
+000000000005afb0 T PyInit__ixformer_torch
+000000000004d870 T _ZN18ixformer_torch_ext12t5_split_qkvERN2at6TensorES2_S2_S2_ll
+0000000000038020 T _ZN18ixformer_torch_ext14ixformer_solveERN2at6TensorES2_b
+000000000003d8e0 T _ZN18ixformer_torch_ext14linear_i8w8o32ERN2at6TensorES2_S2_
+0000000000040a60 T _ZN18ixformer_torch_ext14rms_norm_quantERN2at6TensorES2_S2_d
+000000000003a160 T _ZN18ixformer_torch_ext15ixformer_linearERN2at6TensorES2_RKN3c108optionalIS1_EES7_
+000000000004c530 T _ZN18ixformer_torch_ext15skip_layer_normERN2at6TensorES2_S2_S2_RKN3c108optionalIS1_EES2_bd
+000000000004a650 T _ZN18ixformer_torch_ext16rms_norm_forwardERN2at6TensorES2_S2_d
+0000000000056510 T _ZN18ixformer_torch_ext16vllm_copy_blocksERKSt6vectorIN2at6TensorESaIS2_EES6_RS2_
+0000000000056b00 T _ZN18ixformer_torch_ext16vllm_swap_blocksERN2at6TensorES2_RKSt6vectorIlSaIlEES7_
+0000000000041090 T _ZN18ixformer_torch_ext17vllm_gptq_shuffleERN2at6TensorERKN3c108optionalIS1_EE
+0000000000039ff0 T _ZN18ixformer_torch_ext18get_ipc_shm_tensorERKSt6vectorIlSaIlEEN3c1010ScalarTypeERKNS5_6DeviceEm
+000000000003b1e0 T _ZN18ixformer_torch_ext18ixformer_linear_exERN2at6TensorES2_RKN3c108optionalIS1_EE
+0000000000034550 T _ZN18ixformer_torch_ext18lightllm_glm2_ropeERN2at6TensorES2_S2_
+0000000000049260 T _ZN18ixformer_torch_ext19weight_dequant_gptqERN2at6TensorES2_RKN3c108optionalIS1_EESsi
+000000000003fde0 T _ZN18ixformer_torch_ext20dequant_add_residualERN2at6TensorES2_S2_RKN3c108optionalIS1_EEd
+0000000000033e70 T _ZN18ixformer_torch_ext20gelu_and_mul_forwardERN2at6TensorES2_
+0000000000043820 T _ZN18ixformer_torch_ext20quantized_linear_awqERN2at6TensorES2_S2_RKN3c108optionalIS1_EES7_ii
+000000000004be40 T _ZN18ixformer_torch_ext20silu_and_mul_forwardERN2at6TensorES2_
+0000000000044b00 T _ZN18ixformer_torch_ext21quantized_linear_gptqERN2at6TensorES2_S2_RKN3c108optionalIS1_EES7_ii
+00000000000465c0 T _ZN18ixformer_torch_ext21quantized_linear_int8ERN2at6TensorES2_S2_RKN3c108optionalIS1_EE
+0000000000048bc0 T _ZN18ixformer_torch_ext21weight_dequant_float4ERN2at6TensorES2_Ssii
+0000000000031780 T _ZN18ixformer_torch_ext22geglu_training_forwardERN2at6TensorES2_
+0000000000034ba0 T _ZN18ixformer_torch_ext22lightllm_apply_penaltyERN2at6TensorES2_S2_S2_S2_S2_l
+0000000000031e70 T _ZN18ixformer_torch_ext23geglu_training_backwardERN2at6TensorES2_S2_
+0000000000036190 T _ZN18ixformer_torch_ext23lightllm_tokenattentionERN2at6TensorES2_S2_S2_S2_S2_dllS2_
+0000000000045a40 T _ZN18ixformer_torch_ext23quantized_linear_float4ERN2at6TensorES2_S2_RKN3c108optionalIS1_EEii
+000000000003c410 T _ZN18ixformer_torch_ext25ixformer_linear_allreduceERN2at6TensorES2_RKN3c108optionalIS1_EE
+00000000000474b0 T _ZN18ixformer_torch_ext25ixformer_quantized_linearERN2at6TensorES2_S2_SslRKN3c108optionalIS1_EES7_l
+00000000000504a0 T _ZN18ixformer_torch_ext25tgi_rotary_embedding_neoxERN2at6TensorES2_S2_S1_S2_S1_b
+0000000000040040 T _ZN18ixformer_torch_ext26dequant_silu_and_mul_quantERN2at6TensorES2_ddd
+000000000004b090 T _ZN18ixformer_torch_ext26fused_add_rms_norm_forwardERN2at6TensorES2_S2_dd
+0000000000035930 T _ZN18ixformer_torch_ext26lightllm_destindex_copy_kvERN2at6TensorES2_S2_
+0000000000054f60 T _ZN18ixformer_torch_ext26vllm_rotary_embedding_neoxERN2at6TensorES2_S2_lS2_lb
+0000000000040c10 T _ZN18ixformer_torch_ext27add_residual_rms_norm_quantERN2at6TensorES2_S2_S2_d
+000000000004e530 T _ZN18ixformer_torch_ext28t5_split_qkv_update_kv_cacheERN2at6TensorES2_S2_S2_S2_S2_ll
+00000000000403f0 T _ZN18ixformer_torch_ext29dequant_rotary_embedding_neoxERN2at6TensorES2_S2_lS2_S2_S2_ddb
+00000000000401f0 T _ZN18ixformer_torch_ext30dequant_silu_and_mul_quant_perERN2at6TensorES2_ddS2_S2_
+0000000000055af0 T _ZN18ixformer_torch_ext32vllm_cache_ops_reshape_and_cacheERN2at6TensorES2_S2_S2_S2_ll
+0000000000049ba0 T _ZN18ixformer_torch_ext33ixformer_quantized_weight_dequantERN2at6TensorES2_SsSslRKN3c108optionalIS1_EEl
+0000000000040df0 T _ZN18ixformer_torch_ext35dequant_add_residual_rms_norm_quantERN2at6TensorES2_S2_S2_RKN3c108optionalIS1_EEdd
+00000000000517b0 T _ZN18ixformer_torch_ext37vllm_single_query_cached_kv_attentionERN2at6TensorES2_S2_S2_S2_dS2_S2_lllbRKN3c108optionalIS1_EE
+0000000000053610 T _ZN18ixformer_torch_ext40vllm_single_query_cached_kv_attention_v2ERN2at6TensorElS2_S2_S2_S2_S2_S2_S2_dS2_S2_lllbRKN3c108optionalIS1_EE
+000000000003f720 T _ZN18ixformer_torch_ext5quantERN2at6TensorES2_d
+000000000003fa90 T _ZN18ixformer_torch_ext7dequantERN2at6TensorES2_RKN3c108optionalIS1_EEd
+000000000003f8d0 T _ZN18ixformer_torch_ext9quant_perERN2at6TensorES2_S2_
+0000000000039ea0 T _ZN18ixformer_torch_ext9to_stringERKSt6vectorIlSaIlEE
+0000000000072898 T _fini
+0000000000029000 T _init

cat_files/symbol_dumps/sym_libcuinfer.txt

@@ -0,0 +1,270 @@
+0000000002f30110 T _ZGTtNKSt11logic_error4whatEv
+0000000002f30860 T _ZGTtNKSt13runtime_error4whatEv
+0000000002f2ffa0 T _ZGTtNSt11logic_errorC1EPKc
+0000000002f30030 T _ZGTtNSt11logic_errorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f2ffa0 T _ZGTtNSt11logic_errorC2EPKc
+0000000002f30030 T _ZGTtNSt11logic_errorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f300f0 T _ZGTtNSt11logic_errorD0Ev
+0000000002f300d0 T _ZGTtNSt11logic_errorD1Ev
+0000000002f300d0 T _ZGTtNSt11logic_errorD2Ev
+0000000002f30880 T _ZGTtNSt11range_errorC1EPKc
+0000000002f30910 T _ZGTtNSt11range_errorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30880 T _ZGTtNSt11range_errorC2EPKc
+0000000002f30910 T _ZGTtNSt11range_errorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f309d0 T _ZGTtNSt11range_errorD0Ev
+0000000002f309b0 T _ZGTtNSt11range_errorD1Ev
+0000000002f309b0 T _ZGTtNSt11range_errorD2Ev
+0000000002f30130 T _ZGTtNSt12domain_errorC1EPKc
+0000000002f301c0 T _ZGTtNSt12domain_errorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30130 T _ZGTtNSt12domain_errorC2EPKc
+0000000002f301c0 T _ZGTtNSt12domain_errorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30280 T _ZGTtNSt12domain_errorD0Ev
+0000000002f30260 T _ZGTtNSt12domain_errorD1Ev
+0000000002f30260 T _ZGTtNSt12domain_errorD2Ev
+0000000002f30410 T _ZGTtNSt12length_errorC1EPKc
+0000000002f304a0 T _ZGTtNSt12length_errorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30410 T _ZGTtNSt12length_errorC2EPKc
+0000000002f304a0 T _ZGTtNSt12length_errorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30560 T _ZGTtNSt12length_errorD0Ev
+0000000002f30540 T _ZGTtNSt12length_errorD1Ev
+0000000002f30540 T _ZGTtNSt12length_errorD2Ev
+0000000002f30580 T _ZGTtNSt12out_of_rangeC1EPKc
+0000000002f30610 T _ZGTtNSt12out_of_rangeC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30580 T _ZGTtNSt12out_of_rangeC2EPKc
+0000000002f30610 T _ZGTtNSt12out_of_rangeC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f306d0 T _ZGTtNSt12out_of_rangeD0Ev
+0000000002f306b0 T _ZGTtNSt12out_of_rangeD1Ev
+0000000002f306b0 T _ZGTtNSt12out_of_rangeD2Ev
+0000000002f306f0 T _ZGTtNSt13runtime_errorC1EPKc
+0000000002f30780 T _ZGTtNSt13runtime_errorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f306f0 T _ZGTtNSt13runtime_errorC2EPKc
+0000000002f30780 T _ZGTtNSt13runtime_errorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30840 T _ZGTtNSt13runtime_errorD0Ev
+0000000002f30820 T _ZGTtNSt13runtime_errorD1Ev
+0000000002f30820 T _ZGTtNSt13runtime_errorD2Ev
+0000000002f309f0 T _ZGTtNSt14overflow_errorC1EPKc
+0000000002f30a80 T _ZGTtNSt14overflow_errorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f309f0 T _ZGTtNSt14overflow_errorC2EPKc
+0000000002f30a80 T _ZGTtNSt14overflow_errorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30b40 T _ZGTtNSt14overflow_errorD0Ev
+0000000002f30b20 T _ZGTtNSt14overflow_errorD1Ev
+0000000002f30b20 T _ZGTtNSt14overflow_errorD2Ev
+0000000002f30b60 T _ZGTtNSt15underflow_errorC1EPKc
+0000000002f30bf0 T _ZGTtNSt15underflow_errorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30b60 T _ZGTtNSt15underflow_errorC2EPKc
+0000000002f30bf0 T _ZGTtNSt15underflow_errorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f30cb0 T _ZGTtNSt15underflow_errorD0Ev
+0000000002f30c90 T _ZGTtNSt15underflow_errorD1Ev
+0000000002f30c90 T _ZGTtNSt15underflow_errorD2Ev
+0000000002f302a0 T _ZGTtNSt16invalid_argumentC1EPKc
+0000000002f30330 T _ZGTtNSt16invalid_argumentC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f302a0 T _ZGTtNSt16invalid_argumentC2EPKc
+0000000002f30330 T _ZGTtNSt16invalid_argumentC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+0000000002f303f0 T _ZGTtNSt16invalid_argumentD0Ev
+0000000002f303d0 T _ZGTtNSt16invalid_argumentD1Ev
+0000000002f303d0 T _ZGTtNSt16invalid_argumentD2Ev
+0000000002f2fdf0 T _ZNKSt3_V214error_category10_M_messageEi
+0000000002f2f900 T _ZNSt11logic_errorC1EOS_
+0000000002f2f9f0 T _ZNSt11logic_errorC1EPKc
+0000000002f2f8b0 T _ZNSt11logic_errorC1ERKS_
+0000000002f2f900 T _ZNSt11logic_errorC2EOS_
+0000000002f2f9f0 T _ZNSt11logic_errorC2EPKc
+0000000002f2f8b0 T _ZNSt11logic_errorC2ERKS_
+0000000002f2f930 T _ZNSt11logic_erroraSEOS_
+0000000002f2f8e0 T _ZNSt11logic_erroraSERKS_
+0000000002f2fc50 T _ZNSt11range_errorC1EPKc
+0000000002f2fc50 T _ZNSt11range_errorC2EPKc
+0000000002f2fae0 T _ZNSt12domain_errorC1EPKc
+0000000002f2fae0 T _ZNSt12domain_errorC2EPKc
+0000000002f2fb20 T _ZNSt12length_errorC1EPKc
+0000000002f2fb20 T _ZNSt12length_errorC2EPKc
+0000000002f2fb40 T _ZNSt12out_of_rangeC1EPKc
+0000000002f2fb40 T _ZNSt12out_of_rangeC2EPKc
+0000000002f2f9a0 T _ZNSt13runtime_errorC1EOS_
+0000000002f2fb60 T _ZNSt13runtime_errorC1EPKc
+0000000002f2f950 T _ZNSt13runtime_errorC1ERKS_
+0000000002f2f9a0 T _ZNSt13runtime_errorC2EOS_
+0000000002f2fb60 T _ZNSt13runtime_errorC2EPKc
+0000000002f2f950 T _ZNSt13runtime_errorC2ERKS_
+0000000002f2f9d0 T _ZNSt13runtime_erroraSEOS_
+0000000002f2f980 T _ZNSt13runtime_erroraSERKS_
+0000000002f2fc70 T _ZNSt14overflow_errorC1EPKc
+0000000002f2fc70 T _ZNSt14overflow_errorC2EPKc
+0000000002f2fc90 T _ZNSt15underflow_errorC1EPKc
+0000000002f2fc90 T _ZNSt15underflow_errorC2EPKc
+0000000002f2fb00 T _ZNSt16invalid_argumentC1EPKc
+0000000002f2fb00 T _ZNSt16invalid_argumentC2EPKc
+0000000002f30dd0 T _ZNSt8ios_base7_M_moveERS_
+0000000002f30ee0 T _ZNSt8ios_base7_M_swapERS_
+0000000002f30cd0 T _ZSt24__throw_out_of_range_fmtPKcz
+0000000002f3458c T _fini
+000000000001d000 T _init
+0000000002f28b90 T cuInferPageAttention
+0000000002f28fc0 T cuInferPageAttentionFuse
+0000000002f28560 T cuInferPageAttentionGetWorkspace
+0000000002f28360 T cuInferPageAttentionGetWorkspaceV2
+0000000002f28760 T cuInferPageAttentionV2
+0000000002ef37f0 T cuinferActivationForward
+0000000002f20e10 T cuinferAddTensor
+0000000002ef2790 T cuinferArrangeAttenOutputI8II8O
+0000000002ef2710 T cuinferArrangeEncselfQkvI8II8O
+0000000002ef2dd0 T cuinferArrangeEncselfQkvSepI8II8O
+0000000002ef6680 T cuinferBatchNormalizationForwardInference
+0000000002ef5d80 T cuinferBatchNormalizationForwardTraining
+0000000002ef6ec0 T cuinferBatchNormalizationForwardTrainingEx
+0000000002ef2940 T cuinferBiasGeluI8II8O
+0000000002f218f0 T cuinferBiasResidualLn
+0000000002f0a4b0 T cuinferCTCLoss
+0000000002ef9db0 T cuinferConcatenate
+0000000002f03060 T cuinferConvolutionForward
+0000000002ef2750 T cuinferCorrelationSoftmaxEncselfI32II8O
+0000000002ef2770 T cuinferCorrelationSoftmaxEncselfI8II8O
+0000000002f10870 T cuinferCreate
+0000000002ef3050 T cuinferCreateActivationDescriptor
+0000000002f092e0 T cuinferCreateCTCLossDescriptor
+0000000002efab70 T cuinferCreateConvolutionDescriptor
+0000000002f0c010 T cuinferCreateDropoutDescriptor
+0000000002f0d180 T cuinferCreateFilterDescriptor
+0000000002f0e9d0 T cuinferCreateLRNDescriptor
+0000000002f15660 T cuinferCreatePersistentRNNPlan
+0000000002f11290 T cuinferCreatePoolingDescriptor
+0000000002f15430 T cuinferCreateRNNDescriptor
+0000000002f148a0 T cuinferCreateReduceTensorDescriptor
+0000000002f1f220 T cuinferCreateTensorDescriptor
+0000000002f251b0 T cuinferCropAndResize
+0000000002f21c60 T cuinferCustomGemm
+0000000002f229b0 T cuinferCustomGemmEx
+0000000002f1dcd0 T cuinferDeQuantSoftmaxForwardQuant
+0000000002ef5a20 T cuinferDeriveBNTensorDescriptor
+0000000002f10aa0 T cuinferDestroy
+0000000002ef37c0 T cuinferDestroyActivationDescriptor
+0000000002f0a050 T cuinferDestroyCTCLossDescriptor
+0000000002efc170 T cuinferDestroyConvolutionDescriptor
+0000000002f0c240 T cuinferDestroyDropoutDescriptor
+0000000002f0e1a0 T cuinferDestroyFilterDescriptor
+0000000002f0f4b0 T cuinferDestroyLRNDescriptor
+0000000002f15a70 T cuinferDestroyPersistentRNNPlan
+0000000002f12e30 T cuinferDestroyPoolingDescriptor
+0000000002f15640 T cuinferDestroyRNNDescriptor
+0000000002f20c20 T cuinferDestroyTensorDescriptor
+0000000002f0cab0 T cuinferDropoutForward
+0000000002f0c290 T cuinferDropoutGetReserveSpaceSize
+0000000002f0c270 T cuinferDropoutGetStatesSize
+0000000002ef2600 T cuinferEncEmbI8I
+0000000002ef2670 T cuinferEncEmbI8I_M8I
+0000000002f25420 T cuinferFMHAForward
+0000000002f25c60 T cuinferFMHAForwardEx
+0000000002effd40 T cuinferFindConvolutionForwardAlgorithm
+0000000002f02630 T cuinferFindConvolutionForwardAlgorithmEx
+0000000002f018b0 T cuinferFindConvolutionForwardAlgorithmFP16
+0000000002ef28f0 T cuinferFusedMultiHeadAttentionI8
+0000000002f26340 T cuinferGPTFMHAForward
+0000000002ef3580 T cuinferGetActivationDescriptor
+0000000002ef7f10 T cuinferGetBatchNormalizationForwardTrainingExWorkspaceSize
+0000000002ef7d60 T cuinferGetBatchNormalizationTrainingExReserveSpaceSize
+0000000002f09c00 T cuinferGetCTCLossDescriptor
+0000000002f09e10 T cuinferGetCTCLossDescriptorEx
+0000000002f0a080 T cuinferGetCTCLossWorkspaceSize
+0000000002efbee0 T cuinferGetConvolution2dDescriptor
+0000000002efeec0 T cuinferGetConvolution2dForwardOutputDim
+0000000002eff5b0 T cuinferGetConvolutionForwardAlgorithm
+0000000002f03f60 T cuinferGetConvolutionForwardAlgorithmMaxCount
+0000000002f03dd0 T cuinferGetConvolutionForwardAlgorithm_v7
+0000000002f00700 T cuinferGetConvolutionForwardWorkspaceSize
+0000000002f04020 T cuinferGetConvolutionGroupCount
+0000000002f04030 T cuinferGetConvolutionMathType
+0000000002f04230 T cuinferGetConvolutionNdDescriptor
+0000000002f04540 T cuinferGetConvolutionNdForwardOutputDim
+0000000002f10f60 T cuinferGetCudartVersion
+0000000002f225e0 T cuinferGetCustomGemmExWorkspace
+0000000002f0c870 T cuinferGetDropoutDescriptor
+0000000002f10ed0 T cuinferGetErrorString
+0000000002f0dd90 T cuinferGetFilter4dDescriptor
+0000000002f0df40 T cuinferGetFilterNdDescriptor
+0000000002f20a20 T cuinferGetFilterSizeInBytes
+0000000002f27890 T cuinferGetHammingDistanceWorkspace
+0000000002f0f070 T cuinferGetLRNDescriptor
+0000000002f28270 T cuinferGetNMSBatchedWorkspaceSize
+0000000002f28340 T cuinferGetNMSBatchedYoloFusedWorkspaceSize
+0000000002f281a0 T cuinferGetNMSWorkspaceSize
+0000000002f11b00 T cuinferGetPooling2dDescriptor
+0000000002f12c00 T cuinferGetPooling2dForwardOutputDim
+0000000002f12550 T cuinferGetPoolingNdDescriptor
+0000000002f12940 T cuinferGetPoolingNdForwardOutputDim
+0000000002f10850 T cuinferGetProperty
+0000000002f22f30 T cuinferGetQDEConvolutionTransposedWorkspaceSize
+0000000002f16870 T cuinferGetRNNDescriptor
+0000000002f181e0 T cuinferGetRNNLinLayerBiasParams
+0000000002f17b80 T cuinferGetRNNLinLayerMatrixParams
+0000000002f16dc0 T cuinferGetRNNMatrixMathType
+0000000002f175d0 T cuinferGetRNNParamsSize
+0000000002f166f0 T cuinferGetRNNProjectionLayers
+0000000002f16fc0 T cuinferGetRNNTrainingReserveSize
+0000000002f29dc0 T cuinferGetReduceWorkspace
+0000000002f10d70 T cuinferGetStream
+0000000002f1fce0 T cuinferGetTensor4dDescriptor
+0000000002f20680 T cuinferGetTensorNdDescriptor
+0000000002f20810 T cuinferGetTensorSizeInBytes
+0000000002f2b5e0 T cuinferGetTopKBatchWorkspace
+0000000002f2b370 T cuinferGetTopKWorkspace
+0000000002f10f40 T cuinferGetVersion
+0000000002f271a0 T cuinferGroupNorm
+0000000002f020c0 T cuinferHalfConvolution2dForward
+0000000002f279d0 T cuinferHammingDistance
+0000000002efec30 T cuinferIm2Col
+0000000002f27b60 T cuinferInstanceNorm
+0000000002f0f210 T cuinferLRNCrossChannelForward
+0000000002f10490 T cuinferLSTMForwardInference
+0000000002f27db0 T cuinferLayerNorm
+0000000002ef2c60 T cuinferLayernormResidualI8OFO
+0000000002ef26e0 T cuinferLayernormResualI8O
+0000000002f280f0 T cuinferNMS
+0000000002f281c0 T cuinferNMSBatched
+0000000002f28290 T cuinferNMSBatchedYoloFused
+0000000002f12e60 T cuinferPoolingForward
+0000000002f050c0 T cuinferQConvolutionForward
+0000000002f04ca0 T cuinferQDConvolutionForward
+0000000002f01540 T cuinferQDEConvolutionForward
+0000000002f23790 T cuinferQDEConvolutionTranspose
+0000000002f18840 T cuinferRNNForwardInference
+0000000002f19be0 T cuinferRNNForwardTraining
+0000000002f2a840 T cuinferReduce
+0000000002f14760 T cuinferReduceTensor
+0000000002ef27e0 T cuinferResidualBiasLnI8II8O
+0000000002ef2830 T cuinferResidualBiasLnI8II8OF
+0000000002ef2c30 T cuinferResidualBiaslnI32I
+0000000002ef2aa0 T cuinferResidualBiaslnI32II8O
+0000000002ef27c0 T cuinferResidualBiaslnI8I
+0000000002f15190 T cuinferResize2D
+0000000002f0c6b0 T cuinferRestoreDropoutDescriptor
+0000000002ef32a0 T cuinferSetActivationDescriptor
+0000000002f09530 T cuinferSetCTCLossDescriptor
+0000000002f097f0 T cuinferSetCTCLossDescriptorEx
+0000000002efad00 T cuinferSetConvolution2dDescriptor
+0000000002efb3e0 T cuinferSetConvolutionGroupCount
+0000000002efb690 T cuinferSetConvolutionMathType
+0000000002efb900 T cuinferSetConvolutionNdDescriptor
+0000000002f0c4f0 T cuinferSetDropoutDescriptor
+0000000002f0d3b0 T cuinferSetFilter4dDescriptor
+0000000002f0d920 T cuinferSetFilterNdDescriptor
+0000000002f0ec10 T cuinferSetLRNDescriptor
+0000000002f158c0 T cuinferSetPersistentRNNPlan
+0000000002f114c0 T cuinferSetPooling2dDescriptor
+0000000002f11e10 T cuinferSetPoolingNdDescriptor
+0000000002f15a90 T cuinferSetRNNDescriptor
+0000000002f16b10 T cuinferSetRNNMatrixMathType
+0000000002f163a0 T cuinferSetRNNProjectionLayers
+0000000002f14ae0 T cuinferSetReduceTensorDescriptor
+0000000002f10c10 T cuinferSetStream
+0000000002f1f410 T cuinferSetTensor4dDescriptor
+0000000002f1f980 T cuinferSetTensor4dDescriptorEx
+0000000002f1ff70 T cuinferSetTensorNdDescriptor
+0000000002f20260 T cuinferSetTensorNdDescriptorEx
+0000000002f1d760 T cuinferSoftmaxForward
+0000000002f1ef60 T cuinferSplitForward
+0000000002f2b450 T cuinferTopK
+0000000002f2b650 T cuinferTopKBatch
+0000000002f20c50 T cuinferTransformTensor
+0000000002f2b7d0 T cuinferTranspose
+0000000002ef2870 T cuinferViterbiDecode
+0000000002f2b9e0 T cuinferYoloV5Detect

cat_files/turing_tensorop_gemm.cu

@@ -0,0 +1,354 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without modification, are permitted
+ * provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice, this list of
+ *       conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright notice, this list of
+ *       conditions and the following disclaimer in the documentation and/or other materials
+ *       provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used
+ *       to endorse or promote products derived from this software without specific prior written
+ *       permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR
+ * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
+ * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+ * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+ * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/**
+This example shows how to run matrix multiplication kernels using functions and data structures
+provided by CUTLASS using tensor cores; which we run on a NVIDIA Turing GPU.
+
+Writing a single high performance matrix multiplication kernel is hard but do-able. Whereas writing
+high performance kernels at scale which works for multiple problem sizes with good abstractions is
+really hard. CUTLASS solves this problem by providing simplified abstractions to compose
+multiple sections of gemm kernel. When used properly, the kernels can hit peak performance of GPU
+easily.
+
+CUTLASS divides a kernel into hierarchical composable sections. Which means, at each thread, warp
+and thread-block level, they compute on their own tile-size with higher level of tile sizes being
+composed from lower level ones. Multiple thread-tiles (tile size each thread computes) can be used
+to form warp-tiles (tile size each warp computes) and multiple warp tiles can be used to compute
+threadblock-tile (tile size computed by a threadblock).
+
+In thie example, we split variable initialization into
+1. Setting up data properties : describes how matrices are laid out in the memory and how the kernel
+can view them (logical to physical mapping)
+2. Setting up computation properties : describes how the above set matrices will be used to compute
+output of matrix multiplication.
+
+First, we setup the data types of matrices A, B, C and D along with alpha, beta as the equation for
+GEMM is D = alpha * A * B + beta * C. In CUTLASS, the kernels first compute A * B and leaves the
+rest of the computation to end of the kernel as alpha * X + beta * C is a simple element-wise
+operation on X (A * B) and C. We call this as epilogue of kernel. Hence, we setup data types for
+alpha and beta to be equal to ElementComputeEpilogue = int32_t. As we want to use MMA instructions
+on Turing and they support 8-bit signed integer (int8_t), we use data type for elements in input
+matrix A and B as int8_t. Volta also supports accumulation of partial dot product to int32_t, which
+can store wider range of numbers, we use it as data type of output matrix elements and accumulation.
+We convey this to CUTLASS kernel by initializing template variables ElementAccumulator (int32_t),
+ElementComputeEpilogue (int32_t), ElementInputA (int8_t), ElementInputB (int8_t), ElementOutput
+(int32_t). Communicating just the data type is not enough. As the data is laid out linearly in
+memory, we have to convey the layout of matrices. We do that by initializing template variable
+LayoutInputA to column major cutlass variable, LayoutInputB to row major and LayoutOutput to row
+major. Next, we setup rules to comptue alpha * X + beta * C which is called epilogue of the kernel.
+We initialize template variable EpilogueOp, which takes the data type of output ElementOutput
+(int32_t), the number of elements per vector memory access (16), data type of accumulator (int32_t)
+and data type of computation of linear combination (alpha * X + beta * C).
+
+Now that we setup the properties of data, we have to setup properties of computation.
+
+Second, we create template variables of tile sizes for thread-block, warp and mma-op to 128x256x64,
+64x64x16, 8x8x16 (MxNxK) respectively. When passed to instantiate CUTLASS GEMM kernel, it internally
+deduce the amount of threads needed per thread-block, amount of shared memory, storing data in
+bank-conflict free manner, and ton of other variables required to compose, intialize and launch a
+high performance GEMM kernel. This is the beauty of CUTLASS, it relieves developer from
+understanding and coding complicated hardware optimizations which can easily go wrong.
+
+CUTLASS also supports multiple MMA pipelines in a threadblock. What are MMA pipelines? MMA pipelines
+constitute the whole process of loading input data from global memory to shared memory, loading data
+from shared memory to registers, doing matrix multiplication, store to global memory. The below flow
+sequence shows a typical mma pipeline.
+
+matrix in global memory -> registers -> tile in shared memory -> registers -> mma -> registers ->
+output to global memory
+
+The problem with single pipeline is, each stage is synchronous which means, each stage has to wait
+until the previous finished executing. There are stages in the pipeline which do not have fixed
+latency, for example, the loads from global memory and shared memory. Therefore, we can add one more
+pipeline with a phase shift in mma kernel to hide latency from global and shared memory loads.
+Finally, the pipeline in a kernel looks like
+
+(1) matrix in global memory -> (2) registers -> (3) tile in shared memory -> (4) registers -> (5)
+mma -> (6) registers -> (7) output to global memory (1) <null> -> (2) <null> -> (3) matrix in global
+memory -> (4) registers -> (5) tile in shared memory -> (6) registers -> (7) mma -> (8) registers ->
+(9) output to global memory
+
+This way, you can hide the second global memoroy load latency by doing computation on already loaded
+input data.
+
+There are few more template variables initialized such as, which threadblock tile of output matrix
+is done which threadblock launched on an SM, CUDA SM architecture of GPU you want to run on.
+
+These are all put together to create a template variable which describes CUTLASS GEMM kernel using
+cutlass::gemm::device::Gemm template.
+
+The next step is to intialize physical data, instantiate and initialize CUTLASS kernel and run it.
+We use CUTLASS utilities to initialize, fill, compare matrices as they are simple and doesn't come
+in the way of learning CUTLASS.
+
+Once all the matrices are initialized and filled with data, create arguments tuple to launch CUTLASS
+kernel which takes problem size (M = 5120, N = 4096 and K = 4096), matrices, alpha, beta and the
+important one, split k-dimension factor. Along with that, we query CUTLASS if any scratch-space
+memory required by the kernel we instantiated. If yes, we create it and pass it along with other
+arguments created to intialize CUTLASS kernel then, the kernel is launched.
+
+In this example, we later on launch a reference gemm kernel (from CUTLASS utilities) to compare if
+the output from CUTLASS kernel is same as reference GEMM kernel.
+*/
+
+#include <iostream>
+
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/device/gemm.h"
+#include "cutlass/util/host_tensor.h"
+#include "cutlass/util/reference/device/gemm.h"
+#include "cutlass/util/reference/host/tensor_compare.h"
+#include "cutlass/util/reference/host/tensor_copy.h"
+#include "cutlass/util/reference/host/tensor_fill.h"
+#include "cutlass/util/tensor_view_io.h"
+#include "helper.h"
+
+// The code section below describes datatype for input, output matrices and computation between
+// elements in input matrices.
+using ElementAccumulator = int32_t;                 // <- data type of accumulator
+using ElementComputeEpilogue = ElementAccumulator;  // <- data type of epilogue operations
+using ElementInputA = int8_t;                       // <- data type of elements in input matrix A
+using ElementInputB = int8_t;                       // <- data type of elements in input matrix B
+using ElementOutput = int32_t;                      // <- data type of elements in output matrix D
+
+// The code section below describes matrix layout of input and output matrices. Column Major for
+// Matrix A, Row Major for Matrix B and Row Major for Matrix C
+using LayoutInputA = cutlass::layout::RowMajor;
+using LayoutInputB = cutlass::layout::ColumnMajor;
+using LayoutOutput = cutlass::layout::RowMajor;
+
+// This code section describes whether you want to use tensor cores or regular SIMT cores on GPU SM
+using MMAOp = cutlass::arch::OpClassTensorOp;
+
+// This code section describes CUDA SM architecture number
+using SmArch = cutlass::arch::Sm75;
+
+// This code section describes the tile size a thread block will compute
+using ShapeMMAThreadBlock =
+    cutlass::gemm::GemmShape<128, 256, 64>;  // <- threadblock tile M = 128, N = 256, K = 64
+// This code section describes tile size a warp will compute
+using ShapeMMAWarp = cutlass::gemm::GemmShape<64, 64, 64>;  // <- warp tile M = 64, N = 64, K = 64 
+// This code section describes the size of MMA op
+using ShapeMMAOp = cutlass::gemm::GemmShape<8, 8, 16>;  // <- MMA Op tile M = 8, N = 8, K = 16
+
+// This code section describes how threadblocks are scheduled on GPU
+using SwizzleThreadBlock = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;  // <- ??
+
+// This code section describes the epilogue part of the kernel
+using EpilogueOp = cutlass::epilogue::thread::LinearCombination<
+    ElementOutput,                                     // <- data type of output matrix
+    128 / cutlass::sizeof_bits<ElementOutput>::value,  // <- the number of elements per vectorized
+                                                       // memory access. For a byte, it's 16
+                                                       // elements. This becomes the vector width of
+                                                       // math instructions in the epilogue too
+    ElementAccumulator,                                // <- data type of accumulator
+    ElementComputeEpilogue>;  // <- data type for alpha/beta in linear combination function
+
+// Number of pipelines you want to use
+constexpr int NumStages = 2;
+
+using Gemm = cutlass::gemm::device::Gemm<ElementInputA,
+                                         LayoutInputA,
+                                         ElementInputB,
+                                         LayoutInputB,
+                                         ElementOutput,
+                                         LayoutOutput,
+                                         ElementAccumulator,
+                                         MMAOp,
+                                         SmArch,
+                                         ShapeMMAThreadBlock,
+                                         ShapeMMAWarp,
+                                         ShapeMMAOp,
+                                         EpilogueOp,
+                                         SwizzleThreadBlock,
+                                         NumStages>;
+
+int run() {
+
+  // Turing Tensor Core operations exposed with mma.sync and ldmatrix are first available
+  // in CUDA 10.2. 
+  //
+  // CUTLASS must be compiled with CUDA 10.2 Toolkit to run these examples.
+  if (!(__CUDACC_VER_MAJOR__ > 10 || (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2))) {
+    std::cerr << "Turing Tensor Core operations must be compiled with CUDA 10.2 Toolkit or later." << std::endl;
+    return -1;
+  }
+
+  cudaDeviceProp props;
+
+  cudaError_t error = cudaGetDeviceProperties(&props, 0);
+  if (error != cudaSuccess) {
+    std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
+    return -1;
+  }
+
+  if (!((props.major * 10 + props.minor) >= 75)) {
+    std::cerr << "Turing Tensor Core operations must be run on a machine with compute capability at least 75."
+              << std::endl;
+
+    // Return 0 so tests are considered passing if run on unsupported platforms.
+    return 0;
+  }
+
+  const int length_m = 5120;
+  const int length_n = 4096;
+  const int length_k = 4096;
+
+  // Create a tuple of problem size for matrix multiplication
+  cutlass::gemm::GemmCoord problem_size(length_m, length_n, length_k);
+
+  // Initialize tensors using CUTLASS helper functions
+  cutlass::HostTensor<ElementInputA, LayoutInputA> tensor_a(
+      problem_size.mk());  // <- Create matrix A with dimensions M x K
+  cutlass::HostTensor<ElementInputB, LayoutInputB> tensor_b(
+      problem_size.kn());  // <- Create matrix B with dimensions K x N
+  cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_c(
+      problem_size.mn());  // <- Create matrix C with dimensions M x N
+  cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_d(
+      problem_size.mn());  // <- Create matrix D with dimensions M x N used to store output from
+                           // CUTLASS kernel
+  cutlass::HostTensor<ElementOutput, LayoutOutput> tensor_ref_d(
+      problem_size.mn());  // <- Create matrix D with dimensions M x N used to store output from
+                           // reference kernel
+
+  // Fill input and output matrices on host using CUTLASS helper functions
+  cutlass::reference::host::TensorFillRandomUniform(
+      tensor_a.host_view(),
+      1,
+      ElementInputA(4),
+      ElementInputA(-4),
+      0);  // <- Fill matrix A on host with uniform-distribution random data
+  cutlass::reference::host::TensorFillRandomUniform(
+      tensor_b.host_view(),
+      1,
+      ElementInputB(4),
+      ElementInputB(-4),
+      0);  // <- Fill matrix B on host with uniform-distribution random data
+  cutlass::reference::host::TensorFillRandomUniform(
+      tensor_c.host_view(),
+      1,
+      ElementOutput(4),
+      ElementOutput(-4),
+      0);  // <- Fill matrix C on host with uniform-distribution random data
+  cutlass::reference::host::TensorFill(
+      tensor_d.host_view());  // <- fill matrix D on host with zeros
+  cutlass::reference::host::TensorFill(
+      tensor_ref_d.host_view());  // <- fill matrix D for reference on host with zeros
+
+  // Copy data from host to GPU
+  tensor_a.sync_device();
+  tensor_b.sync_device();
+  tensor_c.sync_device();
+  tensor_d.sync_device();
+  tensor_ref_d.sync_device();
+
+  // Initialize alpha and beta for dot product computation
+  ElementComputeEpilogue alpha = ElementComputeEpilogue(1);
+  ElementComputeEpilogue beta = ElementComputeEpilogue(0);
+
+  // Split K dimension into 1 partitions
+  int split_k_slices = 1;
+
+  // Create a tuple of gemm kernel arguments. This is later passed as arguments to launch
+  // instantiated CUTLASS kernel
+  typename Gemm::Arguments arguments{problem_size,  // <- problem size of matrix multiplication
+                                     tensor_a.device_ref(),  // <- reference to matrix A on device
+                                     tensor_b.device_ref(),  // <- reference to matrix B on device
+                                     tensor_c.device_ref(),  // <- reference to matrix C on device
+                                     tensor_d.device_ref(),  // <- reference to matrix D on device
+                                     {alpha, beta},          // <- tuple of alpha and beta
+                                     split_k_slices};        // <- k-dimension split factor
+
+  // Using the arguments, query for extra workspace required for matrix multiplication computation
+  size_t workspace_size = Gemm::get_workspace_size(arguments);
+
+  // Allocate workspace memory
+  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+
+  // Instantiate CUTLASS kernel depending on templates
+  Gemm gemm_op;
+
+  // Initialize CUTLASS kernel with arguments and workspace pointer
+  cutlass::Status status = gemm_op.initialize(arguments, workspace.get());
+  CUTLASS_CHECK(status);
+
+  // Launch initialized CUTLASS kernel
+  status = gemm_op();
+  CUTLASS_CHECK(status);
+
+  // Create instantiation for device reference gemm kernel
+  cutlass::reference::device::Gemm<ElementInputA,
+                                   LayoutInputA,
+                                   ElementInputB,
+                                   LayoutInputB,
+                                   ElementOutput,
+                                   LayoutOutput,
+                                   ElementComputeEpilogue,
+                                   ElementComputeEpilogue>
+      gemm_device;
+
+  // Launch device reference gemm kernel
+  gemm_device(problem_size,
+              alpha,
+              tensor_a.device_ref(),
+              tensor_b.device_ref(),
+              beta,
+              tensor_c.device_ref(),
+              tensor_ref_d.device_ref());
+
+  // Wait for kernels to finish
+  cudaDeviceSynchronize();
+
+  // Copy output data from CUTLASS and reference kernel to host for comparison
+  tensor_d.sync_host();
+  tensor_ref_d.sync_host();
+
+  // Check if output from CUTLASS kernel and reference kernel are equal or not
+  bool passed = cutlass::reference::host::TensorEquals(
+    tensor_d.host_view(),
+    tensor_ref_d.host_view());
+
+  std::cout << (passed ? "Passed" : "Failed") << std::endl;
+
+  return (passed ? 0  : -1);
+}
+
+int main() {
+  // Turing Tensor Core operations exposed with mma.sync and ldmatrix are first available
+  // in CUDA 10.2. 
+  //
+  // CUTLASS must be compiled with CUDA 10.2 Toolkit to run these examples.
+  if (!(__CUDACC_VER_MAJOR__ > 10 || (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2))) {
+    std::cerr << "Turing Tensor Core operations must be compiled with CUDA 10.2 Toolkit or later." << std::endl;
+
+    // Returning zero so this test passes when built on older Toolkits. 
+    return 0;
+  }
+  else {
+    return run();
+  }
+}
+

cat_ixformer_vllm.py

@@ -0,0 +1,4 @@
+#!/usr/bin/env python3
+"""Print ixformer vllm.py source code."""
+with open("/usr/local/corex/lib64/python3/dist-packages/ixformer/functions/vllm.py") as f:
+    print(f.read())

computility-run.yaml

@@ -8,9 +8,9 @@ command:
     - --served-model-name
     - llm
     - --max-model-len
-    - '80000'
+    - '131072'
     - --gpu-memory-utilization
-    - '0.90'
+    - '0.92'
     - --trust-remote-code
     - -tp
     - '4'
@@ -19,7 +19,7 @@ command:
     - --disable-log-requests
     - --disable-frontend-multiprocessing
     - --max-num-batched-tokens
-    - '8192'
+    - '4096'
     - --enable-chunked-prefill
     - --max-seq-len-to-capture
     - '32768'
@@ -35,8 +35,12 @@ command:
 env:
     - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
       value: 3600
+    - name: BI100_MAX_NUM_SEQS
+      value: 2
     - name: BI100_MOE_COREX_DIRECT_ROUTED
       value: 1
+    - name: BI100_MOE_COREX_TOPK_SOFTMAX
+      value: 1
     - name: BI100_GDN_COREX_PACKED_DECODE
       value: 1
     - name: BI100_HYBRID_KV_ACCOUNTING
@@ -45,7 +49,5 @@ env:
       value: admission64
     - name: BI100_GDN_RESTORE_MODE
       value: hybrid64
-    - name: BI100_MOE_COREX_TOPK_SOFTMAX
-      value: '1'
-    - name: PYTORCH_CUDA_ALLOC_CONF
-      value: expandable_segments:True
+    - name: VLLM_IMAGE_FETCH_TIMEOUT
+      value: 10

core/runtime/py_attention_metadata.cpp

@@ -0,0 +1,348 @@
+/* Adapted from xLLM commit 78aa2a85 (PR #2258).
+   Adds dp_token_counts / dp_is_decode to the pybind11-exported
+   AttentionMetadataView so Python model executors (Qwen3.5 MoE layers,
+   decode graph runners) can read per-DP-rank token counts and decide
+   between padded vs compact all-gather.
+
+   Original: xllm/core/runtime/py_attention_metadata.cpp
+   Scope:    Qwen3.5 data-parallel support in project_6.
+==============================================================================*/
+
+#include "core/runtime/py_attention_metadata.h"
+
+#include <pybind11/stl.h>
+#include <torch/extension.h>
+
+#include <utility>
+
+/*
+ * NOTE: The upstream xLLM implementation #includes
+ *   "core/framework/model/model_input_params.h"
+ *   "core/layers/common/attention_metadata.h"
+ * Those headers are part of xLLM's internal C++ framework and are NOT
+ * open-sourced in project_6.  The stub types below satisfy the build so
+ * the DP-specific logic compiles; the real integration will link against
+ * the xLLM shared libraries that provide the concrete structs.
+ */
+
+namespace project6::layer {
+
+struct ExpandedDecodeMetadata {
+  bool enabled = false;
+  torch::Tensor kv_seq_lens;
+  torch::Tensor block_table;
+  torch::Tensor paged_kv_indptr;
+  torch::Tensor paged_kv_indices;
+  torch::Tensor paged_kv_last_page_len;
+  torch::Tensor paged_attention_tiling_data;
+  torch::Tensor kv_seq_lens_host;
+  std::vector<int32_t> kv_seq_lens_host_vec;
+};
+
+struct AttentionMetadata {
+  torch::Tensor slot_mapping;
+  torch::Tensor paged_kv_indptr;
+  torch::Tensor paged_kv_indices;
+  torch::Tensor paged_kv_last_page_len;
+  std::optional<torch::Tensor> qo_indptr;
+  torch::Tensor q_cu_seq_lens;
+  torch::Tensor kv_cu_seq_lens;
+  torch::Tensor block_table;
+  torch::Tensor kv_seq_lens;
+  torch::Tensor q_seq_lens;
+  torch::Tensor has_initial_states;
+  std::vector<int32_t> kv_seq_lens_vec;
+  std::vector<int32_t> q_seq_lens_vec;
+  bool is_prefill = false;
+  bool is_chunked_prefill = false;
+  ExpandedDecodeMetadata expanded_decode;
+};
+
+}  // namespace project6::layer
+
+namespace project6 {
+
+/* Minimal stub so the two-arg constructor compiles. */
+struct ModelInputParams {
+  struct {
+    std::vector<int32_t> raw_dp_global_token_nums;
+    std::vector<int32_t> dp_global_token_nums;
+    std::vector<int32_t> dp_is_decode;
+  } parallel;
+  struct {
+    torch::Tensor linear_state_indices;
+  } embedding;
+};
+
+namespace py = pybind11;
+
+// ---------------------------------------------------------------------------
+// pybind11 registration
+// ---------------------------------------------------------------------------
+
+void register_attention_metadata_views(py::module_& module) {
+  py::class_<PyExpandedDecodeMetadataView>(module, "ExpandedDecodeMetadataView")
+      .def_property_readonly("enabled", &PyExpandedDecodeMetadataView::enabled)
+      .def_property_readonly("kv_seq_lens",
+                             &PyExpandedDecodeMetadataView::kv_seq_lens)
+      .def_property_readonly("block_table",
+                             &PyExpandedDecodeMetadataView::block_table)
+      .def_property_readonly("paged_kv_indptr",
+                             &PyExpandedDecodeMetadataView::paged_kv_indptr)
+      .def_property_readonly("paged_kv_indices",
+                             &PyExpandedDecodeMetadataView::paged_kv_indices)
+      .def_property_readonly(
+          "paged_kv_last_page_len",
+          &PyExpandedDecodeMetadataView::paged_kv_last_page_len)
+      .def_property_readonly(
+          "paged_attention_tiling_data",
+          &PyExpandedDecodeMetadataView::paged_attention_tiling_data)
+      .def_property_readonly("kv_seq_lens_host",
+                             &PyExpandedDecodeMetadataView::kv_seq_lens_host)
+      .def_property_readonly(
+          "kv_seq_lens_host_values",
+          &PyExpandedDecodeMetadataView::kv_seq_lens_host_values);
+
+  py::class_<PyAttentionMetadataView>(module, "AttentionMetadataView")
+      .def_property_readonly("slot_mapping",
+                             &PyAttentionMetadataView::slot_mapping)
+      .def_property_readonly("paged_kv_indptr",
+                             &PyAttentionMetadataView::paged_kv_indptr)
+      .def_property_readonly("paged_kv_indices",
+                             &PyAttentionMetadataView::paged_kv_indices)
+      .def_property_readonly("paged_kv_last_page_len",
+                             &PyAttentionMetadataView::paged_kv_last_page_len)
+      .def_property_readonly("qo_indptr", &PyAttentionMetadataView::qo_indptr)
+      .def_property_readonly("q_cu_seq_lens",
+                             &PyAttentionMetadataView::q_cu_seq_lens)
+      .def_property_readonly("kv_cu_seq_lens",
+                             &PyAttentionMetadataView::kv_cu_seq_lens)
+      .def_property_readonly("kv_seq_lens_host",
+                             &PyAttentionMetadataView::kv_seq_lens_host)
+      .def_property_readonly("kv_seq_lens_host_values",
+                             &PyAttentionMetadataView::kv_seq_lens_host_values)
+      .def_property_readonly("q_seq_lens_host",
+                             &PyAttentionMetadataView::q_seq_lens_host)
+      .def_property_readonly("block_table",
+                             &PyAttentionMetadataView::block_table)
+      .def_property_readonly("kv_seq_lens",
+                             &PyAttentionMetadataView::kv_seq_lens)
+      .def_property_readonly("linear_state_indices",
+                             &PyAttentionMetadataView::linear_state_indices)
+      .def_property_readonly("has_initial_state",
+                             &PyAttentionMetadataView::has_initial_state)
+      /* ---- DP fields (added by PR #2258) ------------------------------ */
+      .def_property_readonly("dp_token_counts",
+                             &PyAttentionMetadataView::dp_token_counts)
+      .def_property_readonly("dp_is_decode",
+                             &PyAttentionMetadataView::dp_is_decode)
+      /* ----------------------------------------------------------------- */
+      .def_property_readonly("q_seq_lens", &PyAttentionMetadataView::q_seq_lens)
+      .def_property_readonly("expanded_decode_metadata",
+                             &PyAttentionMetadataView::expanded_decode_metadata)
+      .def_property_readonly("is_prefill", &PyAttentionMetadataView::is_prefill)
+      .def_property_readonly("is_chunked_prefill",
+                             &PyAttentionMetadataView::is_chunked_prefill);
+}
+
+// ---------------------------------------------------------------------------
+// PyExpandedDecodeMetadataView
+// ---------------------------------------------------------------------------
+
+PyExpandedDecodeMetadataView::PyExpandedDecodeMetadataView(
+    std::shared_ptr<layer::AttentionMetadata> metadata)
+    : metadata_(std::move(metadata)) {}
+
+bool PyExpandedDecodeMetadataView::enabled() const {
+  return metadata().enabled;
+}
+
+py::object PyExpandedDecodeMetadataView::kv_seq_lens() const {
+  return metadata().kv_seq_lens.defined() ? py::cast(metadata().kv_seq_lens)
+                                          : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::block_table() const {
+  return metadata().block_table.defined() ? py::cast(metadata().block_table)
+                                          : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_kv_indptr() const {
+  return metadata().paged_kv_indptr.defined()
+             ? py::cast(metadata().paged_kv_indptr)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_kv_indices() const {
+  return metadata().paged_kv_indices.defined()
+             ? py::cast(metadata().paged_kv_indices)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_kv_last_page_len() const {
+  return metadata().paged_kv_last_page_len.defined()
+             ? py::cast(metadata().paged_kv_last_page_len)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::paged_attention_tiling_data() const {
+  return metadata().paged_attention_tiling_data.defined()
+             ? py::cast(metadata().paged_attention_tiling_data)
+             : py::none();
+}
+
+py::object PyExpandedDecodeMetadataView::kv_seq_lens_host() const {
+  return metadata().kv_seq_lens_host.defined()
+             ? py::cast(metadata().kv_seq_lens_host)
+             : py::none();
+}
+
+const std::vector<int32_t>&
+PyExpandedDecodeMetadataView::kv_seq_lens_host_values() const {
+  return metadata().kv_seq_lens_host_vec;
+}
+
+const layer::ExpandedDecodeMetadata& PyExpandedDecodeMetadataView::metadata()
+    const {
+  return metadata_->expanded_decode;
+}
+
+// ---------------------------------------------------------------------------
+// PyAttentionMetadataView
+// ---------------------------------------------------------------------------
+
+PyAttentionMetadataView::PyAttentionMetadataView(
+    std::shared_ptr<layer::AttentionMetadata> metadata)
+    : metadata_(std::move(metadata)),
+      kv_seq_lens_host_(
+          make_host_int32_view(metadata_, metadata_->kv_seq_lens_vec)),
+      q_seq_lens_host_(
+          make_host_int32_view(metadata_, metadata_->q_seq_lens_vec)) {}
+
+PyAttentionMetadataView::PyAttentionMetadataView(
+    std::shared_ptr<layer::AttentionMetadata> metadata,
+    const ModelInputParams& params)
+    : PyAttentionMetadataView(std::move(metadata)) {
+  linear_state_indices_ = params.embedding.linear_state_indices;
+
+  /* ---- DP fields (added by PR #2258) ---------------------------------- */
+  dp_token_counts_ = params.parallel.raw_dp_global_token_nums.empty()
+                         ? params.parallel.dp_global_token_nums
+                         : params.parallel.raw_dp_global_token_nums;
+  dp_is_decode_ = params.parallel.dp_is_decode;
+  /* --------------------------------------------------------------------- */
+}
+
+const torch::Tensor& PyAttentionMetadataView::slot_mapping() const {
+  return metadata_->slot_mapping;
+}
+
+const torch::Tensor& PyAttentionMetadataView::paged_kv_indptr() const {
+  return metadata_->paged_kv_indptr;
+}
+
+const torch::Tensor& PyAttentionMetadataView::paged_kv_indices() const {
+  return metadata_->paged_kv_indices;
+}
+
+const torch::Tensor& PyAttentionMetadataView::paged_kv_last_page_len() const {
+  return metadata_->paged_kv_last_page_len;
+}
+
+py::object PyAttentionMetadataView::qo_indptr() const {
+  if (!metadata_->qo_indptr.has_value() || !metadata_->qo_indptr->defined()) {
+    return py::none();
+  }
+  return py::cast(*metadata_->qo_indptr);
+}
+
+py::object PyAttentionMetadataView::q_cu_seq_lens() const {
+  return optional_tensor(metadata_->q_cu_seq_lens);
+}
+
+py::object PyAttentionMetadataView::kv_cu_seq_lens() const {
+  return optional_tensor(metadata_->kv_cu_seq_lens);
+}
+
+py::object PyAttentionMetadataView::kv_seq_lens_host() const {
+  return optional_tensor(kv_seq_lens_host_);
+}
+
+const std::vector<int32_t>& PyAttentionMetadataView::kv_seq_lens_host_values()
+    const {
+  return metadata_->kv_seq_lens_vec;
+}
+
+py::object PyAttentionMetadataView::block_table() const {
+  return optional_tensor(metadata_->block_table);
+}
+
+py::object PyAttentionMetadataView::kv_seq_lens() const {
+  return optional_tensor(metadata_->kv_seq_lens);
+}
+
+py::object PyAttentionMetadataView::linear_state_indices() const {
+  return optional_tensor(linear_state_indices_);
+}
+
+py::object PyAttentionMetadataView::has_initial_state() const {
+  return optional_tensor(metadata_->has_initial_states);
+}
+
+/* ---- DP fields (added by PR #2258) ------------------------------------ */
+const std::vector<int32_t>& PyAttentionMetadataView::dp_token_counts() const {
+  return dp_token_counts_;
+}
+
+const std::vector<int32_t>& PyAttentionMetadataView::dp_is_decode() const {
+  return dp_is_decode_;
+}
+/* ----------------------------------------------------------------------- */
+
+py::object PyAttentionMetadataView::q_seq_lens() const {
+  return optional_tensor(metadata_->q_seq_lens);
+}
+
+py::object PyAttentionMetadataView::q_seq_lens_host() const {
+  return optional_tensor(q_seq_lens_host_);
+}
+
+PyExpandedDecodeMetadataView PyAttentionMetadataView::expanded_decode_metadata()
+    const {
+  return PyExpandedDecodeMetadataView(metadata_);
+}
+
+bool PyAttentionMetadataView::is_prefill() const {
+  return metadata_->is_prefill;
+}
+
+bool PyAttentionMetadataView::is_chunked_prefill() const {
+  return metadata_->is_chunked_prefill;
+}
+
+torch::Tensor PyAttentionMetadataView::make_host_int32_view(
+    const std::shared_ptr<layer::AttentionMetadata>& metadata,
+    std::vector<int32_t>& host_vec) {
+  if (host_vec.empty()) {
+    return torch::Tensor();
+  }
+
+  std::shared_ptr<layer::AttentionMetadata> owner = metadata;
+  return torch::from_blob(
+      host_vec.data(),
+      {static_cast<int64_t>(host_vec.size())},
+      [owner = std::move(owner)](void*) mutable { owner.reset(); },
+      torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU));
+}
+
+py::object PyAttentionMetadataView::optional_tensor(
+    const torch::Tensor& tensor) {
+  return tensor.defined() ? py::cast(tensor) : py::none();
+}
+
+}  // namespace project6
+
+PYBIND11_MODULE(py_attention_metadata, m) {
+  m.doc() = "DP-aware attention metadata (project6, ported from xLLM PR #2258)";
+  project6::register_attention_metadata_views(m);
+}

core/runtime/py_attention_metadata.h

@@ -0,0 +1,100 @@
+/* Adapted from xLLM commit 78aa2a85 (PR #2258).
+   Adds dp_token_counts / dp_is_decode fields to PyAttentionMetadataView
+   so the Python attention backend can partition KV cache by DP group.
+
+   Original: xllm/core/runtime/py_attention_metadata.h
+   Scope:    Qwen3.5 data-parallel support in project_6.
+==============================================================================*/
+
+#pragma once
+
+#include <pybind11/pybind11.h>
+#include <torch/torch.h>
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+/* Forward declarations — project_6 keeps these in its own layer namespace. */
+namespace project6::layer {
+struct AttentionMetadata;
+struct ExpandedDecodeMetadata;
+}  // namespace project6::layer
+
+namespace project6 {
+
+struct ModelInputParams;
+
+void register_attention_metadata_views(pybind11::module_& module);
+
+class PyExpandedDecodeMetadataView final {
+ public:
+  explicit PyExpandedDecodeMetadataView(
+      std::shared_ptr<layer::AttentionMetadata> metadata);
+
+  bool enabled() const;
+  pybind11::object kv_seq_lens() const;
+  pybind11::object block_table() const;
+  pybind11::object paged_kv_indptr() const;
+  pybind11::object paged_kv_indices() const;
+  pybind11::object paged_kv_last_page_len() const;
+  pybind11::object paged_attention_tiling_data() const;
+  pybind11::object kv_seq_lens_host() const;
+  const std::vector<int32_t>& kv_seq_lens_host_values() const;
+
+ private:
+  const layer::ExpandedDecodeMetadata& metadata() const;
+
+  std::shared_ptr<layer::AttentionMetadata> metadata_;
+};
+
+class PyAttentionMetadataView final {
+ public:
+  explicit PyAttentionMetadataView(
+      std::shared_ptr<layer::AttentionMetadata> metadata);
+  PyAttentionMetadataView(std::shared_ptr<layer::AttentionMetadata> metadata,
+                          const ModelInputParams& params);
+
+  const torch::Tensor& slot_mapping() const;
+  const torch::Tensor& paged_kv_indptr() const;
+  const torch::Tensor& paged_kv_indices() const;
+  const torch::Tensor& paged_kv_last_page_len() const;
+  pybind11::object qo_indptr() const;
+  pybind11::object q_cu_seq_lens() const;
+  pybind11::object kv_cu_seq_lens() const;
+  pybind11::object kv_seq_lens_host() const;
+  const std::vector<int32_t>& kv_seq_lens_host_values() const;
+  pybind11::object q_seq_lens_host() const;
+  pybind11::object block_table() const;
+  pybind11::object kv_seq_lens() const;
+  pybind11::object linear_state_indices() const;
+  pybind11::object has_initial_state() const;
+
+  /* ---- DP fields (added by PR #2258) ---------------------------------- */
+  const std::vector<int32_t>& dp_token_counts() const;
+  const std::vector<int32_t>& dp_is_decode() const;
+  /* --------------------------------------------------------------------- */
+
+  pybind11::object q_seq_lens() const;
+  PyExpandedDecodeMetadataView expanded_decode_metadata() const;
+  bool is_prefill() const;
+  bool is_chunked_prefill() const;
+
+ private:
+  static torch::Tensor make_host_int32_view(
+      const std::shared_ptr<layer::AttentionMetadata>& metadata,
+      std::vector<int32_t>& host_vec);
+  static pybind11::object optional_tensor(const torch::Tensor& tensor);
+
+  std::shared_ptr<layer::AttentionMetadata> metadata_;
+  torch::Tensor kv_seq_lens_host_;
+  torch::Tensor q_seq_lens_host_;
+  torch::Tensor linear_state_indices_;
+
+  /* ---- DP fields (added by PR #2258) ---------------------------------- */
+  std::vector<int32_t> dp_token_counts_;
+  std::vector<int32_t> dp_is_decode_;
+  /* --------------------------------------------------------------------- */
+};
+
+}  // namespace project6

debug_gdn_nan.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+"""Debug NaN in C++ torch_chunk_gated_delta_rule.
+
+Tests with smaller dimensions to isolate the issue.
+"""
+import sys
+import os
+import importlib.util
+import torch
+
+def load_mod():
+    so = "/tmp/gdn_test/corex_gdn_chunk_recurrent.so"
+    if not os.path.exists(so):
+        print("Run verify_gdn_cpp.py first to compile")
+        return None
+    spec = importlib.util.spec_from_file_location("corex_gdn_chunk_recurrent", so)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+def main():
+    mod = load_mod()
+    if mod is None:
+        return 1
+
+    # Test with tiny dimensions to isolate
+    for T in [1, 2, 4, 8, 16, 32, 64, 128]:
+        torch.manual_seed(42)
+        B = 1
+        Hk, Hv, D = 4, 8, 128
+        chunk = min(64, T)
+
+        q = torch.randn(B, T, Hk, D, device="cuda", dtype=torch.float16)
+        k = torch.randn(B, T, Hk, D, device="cuda", dtype=torch.float16)
+        v = torch.randn(B, T, Hv, D, device="cuda", dtype=torch.float16)
+        g = torch.randn(B, T, Hv, device="cuda", dtype=torch.float16)
+        beta = torch.randn(B, T, Hv, device="cuda", dtype=torch.float16)
+
+        out, state = mod.torch_chunk_gated_delta_rule(
+            q, k, v, g, beta, chunk, None, True, True)
+
+        has_nan = out.isnan().any().item()
+        nan_count = out.isnan().sum().item() if has_nan else 0
+        print(f"T={T:4d} chunk={chunk:3d}: NaN={has_nan} (count={nan_count}/{out.numel()})")
+
+        if has_nan and T <= 16:
+            # Print where NaN is
+            nan_mask = out.isnan()
+            print(f"  NaN positions: {nan_mask.nonzero()[:5].tolist()}")
+
+    # Test: does chunk_size=T (no actual chunking) work?
+    print("\n--- Single chunk (chunk_size == T) ---")
+    for T in [32, 64]:
+        torch.manual_seed(42)
+        q = torch.randn(1, T, 4, 128, device="cuda", dtype=torch.float16)
+        k = torch.randn(1, T, 4, 128, device="cuda", dtype=torch.float16)
+        v = torch.randn(1, T, 8, 128, device="cuda", dtype=torch.float16)
+        g = torch.randn(1, T, 8, device="cuda", dtype=torch.float16)
+        beta = torch.randn(1, T, 8, device="cuda", dtype=torch.float16)
+
+        out, state = mod.torch_chunk_gated_delta_rule(
+            q, k, v, g, beta, T, None, True, True)
+        print(f"T={T} chunk={T}: NaN={out.isnan().any().item()}")
+
+    # Test: float32 input instead of float16
+    print("\n--- Float32 input ---")
+    for T in [64, 128]:
+        torch.manual_seed(42)
+        q = torch.randn(1, T, 4, 128, device="cuda", dtype=torch.float32)
+        k = torch.randn(1, T, 4, 128, device="cuda", dtype=torch.float32)
+        v = torch.randn(1, T, 8, 128, device="cuda", dtype=torch.float32)
+        g = torch.randn(1, T, 8, device="cuda", dtype=torch.float32)
+        beta = torch.randn(1, T, 8, device="cuda", dtype=torch.float32)
+
+        out, state = mod.torch_chunk_gated_delta_rule(
+            q, k, v, g, beta, 64, None, True, True)
+        print(f"T={T} chunk=64 f32: NaN={out.isnan().any().item()}")
+
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

docs/PORTING_ASSESSMENT.md

@@ -0,0 +1,269 @@
+# BI-V100 移植评估：全仓库编译目标清单
+
+## 架构差异
+
+| | NVIDIA V100 | Iluvatar BI-V100 |
+|---|---|---|
+| 架构标识 | `sm_70` | `ivcore10` |
+| 编译器 | `nvcc` / `clang --cuda-gpu-arch=sm_70` | `corex clang/16 --cuda-gpu-arch=ivcore10` |
+| 运行时编译 | `nvrtc` + `nvjitlink` | **不支持** |
+| Driver API | `cuLibraryLoadData` / `cuLibraryGetKernel` | **不支持** |
+| Tensor Core | HMMA (SM70) | **不支持** |
+| Warp size | 32 | 32 (确认) |
+| SMEM | 96KB (configurable) | 48KB |
+| L2 Cache | 6MB | 不同 |
+| SMs | 80 | 16 |
+| CUB block-level | ✅ header-only | ✅ 可通过 corex clang 编译 |
+| CUB device-level | ✅ via nvrtc JIT | ❌ 需要 AOT 替代方案 |
+
+## 1. NVIDIA/CCCL (10,083 files)
+
+### 1.1 c/parallel SHARED LIBRARY — cccl.c.parallel.so
+
+**状态: ❌ 不能直接移植**
+
+12 个算法全部依赖 NVRTC JIT 编译。每个 .cu 通过 `nvrtc_translation_unit` 生成源码，`-arch=sm_XX` 编译，`cuLibraryLoadData` 加载。
+
+| 算法 | 源文件 | 行数 | NVRTC 依赖 | 移植方案 |
+|---|---|---|---|---|
+| reduce | reduce.cu | 783 | nvrtc × 30 | AOT: 直接调用 cub::DeviceReduce with corex |
+| scan | scan.cu | 943 | nvrtc × 25 | AOT: cub::DeviceScan |
+| radix_sort | radix_sort.cu | 947 | nvrtc × 24 | AOT: cub::DeviceRadixSort |
+| merge_sort | merge_sort.cu | 763 | nvrtc × 25 | AOT: cub::DeviceMergeSort |
+| transform | transform.cu | 1014 | nvrtc × 38 | AOT: cub::DeviceTransform |
+| select_if | three_way_partition.cu | 697 | nvrtc × 29 | AOT: cub::DeviceSelect |
+| histogram | histogram.cu | 858 | nvrtc × 18 | AOT: cub::DeviceHistogram |
+| segmented_reduce | segmented_reduce.cu | 655 | nvrtc × 26 | AOT: cub::DeviceSegmentedReduce |
+| segmented_sort | segmented_sort.cu | 1306 | nvrtc × 40 | AOT: cub::DeviceSegmentedSort |
+| binary_search | binary_search.cu | 547 | nvrtc × 8 | AOT: cub::DeviceBinarySearch |
+| unique_by_key | unique_by_key.cu | 768 | nvrtc × 19 | AOT: cub::DeviceUniqueByKey |
+| for | for.cu | 426 | nvrtc × 15 | AOT: cub::DeviceFor |
+
+**移植策略**: 不搬 c/parallel，而是直接用 CUB header-only API 写 AOT .cu 文件，用 corex clang 编译成 .so。每个算法 = 一组固定类型特化。
+
+### 1.2 c/parallel.v2 SHARED LIBRARY
+
+**状态: ❌ 不能直接移植 (依赖 hostjit/libnvcc)**
+
+v2 用嵌入式 clang 做 JIT，不用 nvrtc。理论上可以用 corex clang 替换 libnvcc 的 clang，但改造量大。
+
+### 1.3 CUB block/warp/thread 原语 (header-only)
+
+**状态: ✅ 可直接使用**
+
+| 类别 | 文件数 | 说明 |
+|---|---|---|
+| block primitives | 25 .cuh | BlockReduce, BlockScan, BlockSort, BlockLoad, BlockStore 等 |
+| warp primitives | 17 .cuh | WarpReduce, WarpScan, WarpSort 等 |
+| thread primitives | 8 .cuh | ThreadReduce, ThreadScan, ThreadSort 等 |
+| agent implementations | 26 .cuh | 每个 device algorithm 的 kernel 实现 |
+| dispatch kernels | 17 .cuh | kernel launch 模板 |
+| tuning policies | 27 .cuh | SM-specific 参数选择 (需适配 ivcore10) |
+
+**移植策略**: `#include <cub/block/block_reduce.cuh>` 直接在 corex .cu 中使用。tuning policy 需要为 ivcore10 写新的参数表。
+
+### 1.4 CUB/Thrust benchmarks + examples
+
+| 类别 | 数量 | 移植状态 |
+|---|---|---|
+| CUB benchmarks | 82 | 需适配 ivcore10 编译 |
+| CUB examples | 18 | 需适配 ivcore10 编译 |
+| Thrust examples | 60 | 需适配 ivcore10 编译 |
+| Thrust benchmarks | 75 | 需适配 ivcore10 编译 |
+| cudax examples | 68 | 依赖 cudax runtime，暂不移植 |
+| libcudacxx benchmarks | 62 | 需适配 ivcore10 编译 |
+
+---
+
+## 2. NVIDIA/CUTLASS (7,787 files)
+
+### 2.1 核心 GEMM 库 (header-only)
+
+**状态: ⚠️ 部分可移植**
+
+| SM 架构 | 文件数 | BI-V100 兼容 |
+|---|---|---|
+| SM70 (Volta SIMT) | ~20 | ✅ 需验证 ivcore10 兼容性 |
+| SM75 (Turing) | ~30 | ⚠️ 部分 (SIMT mode) |
+| SM80 (Ampere Tensor) | ~200 | ❌ 需要 HMMA |
+| SM90 (Hopper) | ~300 | ❌ |
+| SM100/120 (Blackwell) | ~200 | ❌ |
+
+### 2.2 Grouped GEMM (MoE 核心)
+
+| Example | 文件 | SM 要求 | 移植状态 |
+|---|---|---|---|
+| 24_gemm_grouped | gemm_grouped.cu | SM70+ SIMT | ✅ 可移植 |
+| 57_hopper_grouped_gemm | — | SM90 | ❌ |
+| 64_ada_fp8_gemm_grouped | — | SM89 | ❌ |
+| 92_blackwell_moe_gemm | — | SM100 | ❌ |
+
+**移植策略**: example 24 (SIMT grouped GEMM) 是唯一能在 BI-V100 跑的。搬过来，接口适配到 xllm group_gemm。
+
+### 2.3 编译目标汇总
+
+| 类别 | 数量 |
+|---|---|
+| Example executables | 164 .cu |
+| Test executables | 862 .cu |
+| Include headers | 785 |
+| SM70 兼容子集 | ~20 examples + ~50 tests |
+
+---
+
+## 3. Dao-AILab/flash-attention (606 .cu files)
+
+### 3.1 flash_attn_2_cuda.so
+
+**状态: ❌ 不能直接移植 (SM80+ Tensor Core)**
+
+所有 kernel 使用 `cute::MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>` — 依赖 Ampere Tensor Core。
+
+| Kernel 类别 | .cu 数量 | SM 要求 |
+|---|---|---|
+| SM80 fwd | 48 | ❌ Tensor Core |
+| SM80 bwd | 24 | ❌ Tensor Core |
+| SM80 fwd_split | 48 | ❌ Tensor Core |
+| SM80 fwd_split_align | 42 | ❌ Tensor Core |
+| Hopper (SM90+) | 453 | ❌ |
+
+### 3.2 可用的算法模板
+
+| 文件 | 行数 | 价值 |
+|---|---|---|
+| flash_fwd_kernel.h | 1301 | attention 算法流程 (Q×K softmax V) |
+| softmax.h | 189 | online softmax 实现 |
+| kernel_traits.h | 344 | SMEM/register 分配策略 |
+| mask.h | 214 | causal mask 实现 |
+| rotary.h | 153 | RoPE in-kernel 实现 |
+
+**移植策略**: 不搬 .cu kernel（依赖 Tensor Core），搬算法模板头文件，基于 CUB block primitives 重写 SIMT attention kernel for ivcore10。或者直接用 ixformer base image 的 `ixinfer_flash_attn_unpad_with_block_tables`（已编译好）。
+
+### 3.3 Layer Norm kernels
+
+| 类别 | .cu 数量 | SM 要求 |
+|---|---|---|
+| ln_fwd | 14 (256~8192 width) | ✅ 纯 SIMT |
+| ln_bwd | 14 | ✅ 纯 SIMT |
+| ln_parallel_fwd | 14 | ✅ 纯 SIMT |
+| ln_parallel_bwd | 14 | ✅ 纯 SIMT |
+
+**移植策略**: Layer norm kernel 是纯 SIMT，不依赖 Tensor Core。可直接用 corex clang 编译。hidden_size=5120 对应 ln_fwd_5120.cu。
+
+---
+
+## 4. jd-opensource/xllm (全平台推理引擎)
+
+### 4.1 ILU (BI-V100) 专用代码
+
+**状态: ✅ 已在项目中 (upstream_ref + ex_engine)**
+
+| 文件 | 行数 | 作用 | 状态 |
+|---|---|---|---|
+| ilu/activation.cpp | 32 | silu_and_mul → ixformer::infer | ✅ 已搬 |
+| ilu/norm.cpp | 50 | rms_norm → ixformer::infer | ✅ 已搬 |
+| ilu/rope.cpp | 31 | rotary_embedding → ixformer::infer | ✅ 已搬 |
+| ilu/attention.cpp | 162 | prefill + decode → ixformer::infer | ✅ 已搬 |
+| ilu/fused_moe.cpp | 99 | topk + expand + combine → ixformer::infer | ✅ 已搬 |
+| ilu/group_gemm.cpp | 39 | group_gemm → ixformer::infer | ✅ 已搬 |
+| ilu/matmul.cpp | 73 | linear → ixformer::infer | ✅ 已搬 |
+| ilu/ixformer.h | 147 | 完整 ixformer::infer API 声明 | ✅ 已搬 |
+| ilu/ilu_ops_api.h | 153 | xllm kernel 层 API | ✅ 已搬 |
+| ilu/utils.h | 62 | 工具函数 | ✅ 已搬 |
+| layers/ilu/fused_moe.cpp | 806 | 完整 MoE 7步 pipeline | ✅ 已搬 |
+| layers/ilu/attention.cpp | 189 | attention layer 封装 | ✅ 已搬 |
+
+### 4.2 CUDA kernels (SM-agnostic)
+
+| 文件 | 行数 | SM 限制 | 状态 |
+|---|---|---|---|
+| activation.cu | 188 | 无 | ✅ 已搬 |
+| norm.cu | 600 | 需 cub::BlockReduce | ✅ 已搬 |
+| rope.cu | 258 | 无 | ✅ 已搬 |
+| block_copy.cu | 209 | 无 | ✅ 已搬 |
+| reshape_paged_cache.cu | 101 | 无 | ✅ 已搬 |
+| moe/moe_topk_softmax_kernels.cuh | 867 | 无 | ✅ 已搬 |
+| moe/moe_compute_index.cu | 155 | 无 | ✅ 已搬 |
+| moe/moe_combine.cu | 105 | 无 | ✅ 已搬 |
+| moe/moe_fused_topk.cu | 59 | 无 | ✅ 已搬 |
+
+### 4.3 CUDA kernels (SM80+ only)
+
+| 文件 | 行数 | SM 限制 | 移植方案 |
+|---|---|---|---|
+| fused_qknorm_rope.cu | 473 | SM80 (`__CUDA_ARCH__ >= 800`) | 拆出 SIMT 部分 |
+| fp8_quant_utils.cuh | 239 | SM89 (`__CUDA_ARCH__ >= 890`) | 不适用 |
+| cutlass_w8a8/*.cu | ~400 | SM90/100/120 | 不适用 |
+
+### 4.4 其他平台代码 (参考用)
+
+| 平台 | kernel 文件数 | layer 文件数 | 说明 |
+|---|---|---|---|
+| DCU (AMD ROCm) | 14 | 12 | GDN 完整实现可参考 |
+| MLU (Cambricon) | 21 | 35 | GDN + MoE 最完整 |
+| MUSA (Moore Threads) | 14 | 12 | GDN kernel 最近代 |
+| NPU (Ascend) | 30+ | 30+ | tilelang GDN 可参考 |
+
+---
+
+## 5. fla-org/flash-linear-attention (349 Triton kernels)
+
+### 5.1 GatedDeltaNet 专用 kernels
+
+**状态: ⚠️ 需验证 Triton 在 BI-V100 上是否工作**
+
+| 文件 | @triton.jit | 行数 | 说明 |
+|---|---|---|---|
+| chunk_fwd.py | 2 | 428 | GDN 前向 chunk (核心) |
+| fused_recurrent.py | 2 | 478 | GDN decode (单步) |
+| wy_fast.py | 4 | 351 | WY representation |
+| gate.py | 6 | 344 | gate cumsum |
+
+### 5.2 通用 Triton 算子
+
+| 目录 | kernel 数 | 说明 |
+|---|---|---|
+| common/ | 36 | chunk_h, chunk_o, fused_recurrent (所有 linear attention 共享) |
+| utils/ | 44 | cumsum, softmax, matmul, solve_tril |
+| gated_delta_rule/ | 14 | GDN 专用 |
+| gdn2/ | 12 | GDN v2 (新版) |
+| kda/ | 24 | Key-dependent attention |
+| delta_rule/ | 12 | 原始 delta rule |
+| gla/ | 18 | Gated Linear Attention |
+
+### 5.3 Backend 分发
+
+| Backend | SM 要求 | 说明 |
+|---|---|---|
+| FlashQLA | SM90+ | ❌ 不适用 BI-V100 |
+| Triton (default) | 任意 GPU | ⚠️ 需验证 corex Triton |
+| triton_ascend | Ascend NPU | ❌ 不适用 |
+
+---
+
+## 移植优先级
+
+### P0 — 直接可编译 (corex clang ivcore10)
+
+1. **xllm CUDA kernels** (9 files, 2542 lines) — 已搬，需在真机编译测试
+2. **CUB block/warp headers** — 已在 cccl_upstream/，可直接 #include
+3. **ix_moe_bridge.so + ix_attn_bridge.so** — pybind11 桥接 ixformer::infer
+
+### P1 — 需适配后可编 (改 SM 架构 + tuning 参数)
+
+4. **FlashAttention layer_norm kernels** (56 .cu) — 纯 SIMT，改编译 flag
+5. **CUTLASS SM70 SIMT GEMM** (example 24 grouped_gemm) — MoE group_gemm 替代方案
+6. **CUB tuning policies** (27 .cuh) — 为 ivcore10 写参数表 (SMEM=48KB, SM=16)
+
+### P2 — 需要重写 (算法可用，硬件指令不兼容)
+
+7. **FlashAttention fwd kernel** — 基于算法模板用 CUB BlockReduce 重写 SIMT 版
+8. **CCCL c/parallel AOT 版** — 绕过 NVRTC，直接用 CUB device API + corex 编译
+9. **FLA Triton GDN kernels** — 需验证 Triton on corex 可行性
+
+### P3 — 不移植
+
+10. FlashAttention SM80+ Tensor Core kernels
+11. CUTLASS SM80/90/100/120 kernels
+12. CCCL nvrtc/nvjitlink 依赖代码
+13. xllm fp8/cutlass_w8a8 quantization kernels

ex_engine/build_cuinfer_gemm.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+# build_cuinfer_gemm.sh — Compile cuinfer GEMM wrapper
+#
+# Links: libcuinfer.so (from /usr/local/corex/lib64/)
+# Output: cuinfer_gemm_wrapper.so
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+SRC="${SCRIPT_DIR}/cuinfer_gemm_wrapper.cu"
+HDR="${SCRIPT_DIR}/cuinfer_handle.h"
+
+echo "[cuinfer_gemm] Building cuinfer_gemm_wrapper.so"
+
+COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
+CUINFER_LIB=""
+for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib"; do
+    if [[ -f "${d}/libcuinfer.so" ]]; then
+        CUINFER_LIB="${d}"
+        break
+    fi
+done
+
+python3 << PYEOF
+import os, sys, shutil
+
+src = "${SRC}"
+hdr_dir = "${SCRIPT_DIR}"
+cuinfer_lib = "${CUINFER_LIB}"
+
+ldflags = []
+if cuinfer_lib:
+    ldflags = [f"-L{cuinfer_lib}", "-lcuinfer", f"-Wl,-rpath,{cuinfer_lib}"]
+
+try:
+    from torch.utils.cpp_extension import load
+    mod = load(
+        name="cuinfer_gemm_wrapper",
+        sources=[src],
+        extra_include_paths=[hdr_dir],
+        extra_cflags=["-O2", "-std=c++17"],
+        extra_cuda_cflags=["-O2"],
+        extra_ldflags=ldflags,
+        verbose=True,
+    )
+    print("[cuinfer_gemm] ✓ OK")
+
+    import importlib
+    spec = importlib.util.find_spec("cuinfer_gemm_wrapper")
+    if spec and spec.origin:
+        shutil.copy2(spec.origin, os.path.join(hdr_dir, "cuinfer_gemm_wrapper.so"))
+        print(f"[cuinfer_gemm] ✓ Saved")
+
+except Exception as e:
+    print(f"[cuinfer_gemm] ERROR: {e}", file=sys.stderr)
+    sys.exit(1)
+PYEOF

ex_engine/build_gemm_grouped.sh

@@ -0,0 +1,80 @@
+#!/usr/bin/env bash
+# build_gemm_grouped.sh — Compile grouped GEMM kernel + bindings
+#
+# Requires: corex clang/16 + cutlass headers (on BI-V100 device)
+# Output: gemm_grouped.so (importable from Python)
+#
+# Reference: ex_engine/xllm_kernels/build_test_cutlass_batched.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Source files
+GEMM_CU="${SCRIPT_DIR}/csrc/gemm_grouped.cu"
+BIND_CPP="${SCRIPT_DIR}/csrc/gemm_grouped_bind.cpp"
+BATCHED_CU="${SCRIPT_DIR}/../xllm_kernels/cuda/corex_batched_gemm_kernel.cu"
+
+echo "[gemm] Building gemm_grouped.so"
+
+# Find cutlass include path
+SAMPLES="/usr/local/corex-samples-3.2.3_x86_64/samples/cutlass"
+CUTLASS_INCLUDE=""
+for d in "${SAMPLES}/include" "/usr/local/corex/include/cutlass" "/usr/include/cutlass"; do
+    if [[ -d "$d" ]]; then
+        CUTLASS_INCLUDE="$d"
+        break
+    fi
+done
+
+if [[ -z "$CUTLASS_INCLUDE" ]]; then
+    echo "[gemm] ERROR: cutlass include not found"
+    exit 1
+fi
+echo "[gemm] cutlass: ${CUTLASS_INCLUDE}"
+
+python3 << PYEOF
+import os, sys, shutil
+
+script_dir = "${SCRIPT_DIR}"
+cutlass_inc = "${CUTLASS_INCLUDE}"
+
+sources = [
+    "${GEMM_CU}",
+    "${BIND_CPP}",
+    "${BATCHED_CU}",
+]
+sources = [s for s in sources if os.path.isfile(s)]
+
+print(f"[gemm] Compiling {len(sources)} source files")
+for s in sources:
+    print(f"  {os.path.basename(s)}")
+
+try:
+    from torch.utils.cpp_extension import load
+    mod = load(
+        name="gemm_grouped",
+        sources=sources,
+        extra_include_paths=[cutlass_inc, script_dir],
+        extra_cflags=["-O2", "-std=c++17"],
+        extra_ldflags=["/usr/local/corex/lib64/libcuinfer.so", "-Wl,-rpath,/usr/local/corex/lib64"],
+        extra_cuda_cflags=["-O2", "",
+                           f"-I{cutlass_inc}"],
+        verbose=True,
+    )
+    print("[gemm] ✓ Compilation successful")
+
+    import importlib
+    spec = importlib.util.find_spec("gemm_grouped")
+    if spec and spec.origin:
+        dst = os.path.join(script_dir, "gemm_grouped.so")
+        shutil.copy2(spec.origin, dst)
+        print(f"[gemm] ✓ Saved to {dst}")
+
+except Exception as e:
+    print(f"[gemm] ERROR: {e}", file=sys.stderr)
+    import traceback; traceback.print_exc()
+    sys.exit(1)
+PYEOF
+
+echo "[gemm] Done"

ex_engine/build_ix_bridge.sh

@@ -0,0 +1,121 @@
+#!/usr/bin/env bash
+# build_ix_bridge.sh — Compile ix_full_bridge_v2.cpp on BI-V100
+#
+# Upstream ref: xllm/core/kernels/ilu/ixformer.h (all 14 C++ functions)
+# Bridge ref:   ex_engine/csrc/ix_full_bridge_v2.cpp
+#
+# This produces ix_full_bridge_v2.so — a pybind11 module that exposes
+# ALL ixformer::infer functions to Python without any Python fallbacks.
+#
+# Usage:
+#   bash build_ix_bridge.sh [VLLM_ROOT]
+#
+# The .so is deployed to $VLLM_ROOT/ex_engine/ and also to
+# ex_engine/prebuilt/ for the prebuilt pipeline.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+CSRC_DIR="${SCRIPT_DIR}/csrc"
+VLLM_ROOT="${1:-}"
+
+# --- Locate tools ---
+COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
+CLANGXX="${COREX_ROOT}/bin/clang++"
+if [[ ! -x "$CLANGXX" ]]; then
+    CLANGXX=$(command -v clang++ 2>/dev/null || true)
+fi
+if [[ -z "$CLANGXX" ]]; then
+    echo "[ix_bridge] ERROR: clang++ not found" >&2
+    exit 1
+fi
+
+# --- Locate torch and python ---
+PYTHON="${PYTHON:-python3}"
+TORCH_DIR=$($PYTHON -c "import torch; print(torch.utils.cmake_prefix_path)" 2>/dev/null || \
+            $PYTHON -c "import torch; import os; print(os.path.join(os.path.dirname(torch.__file__), 'share', 'cmake'))" 2>/dev/null || true)
+TORCH_INC=$($PYTHON -c "from torch.utils.cpp_extension import include_paths; print(' '.join(['-I'+p for p in include_paths()]))")
+TORCH_LIB=$($PYTHON -c "from torch.utils.cpp_extension import library_paths; print(' '.join(['-L'+p for p in library_paths()]))")
+PYTHON_INC=$($PYTHON -c "from sysconfig import get_paths; print('-I' + get_paths()['include'])")
+
+# --- Locate ixformer .so files for linking ---
+IX_LIBS=""
+for sopath in \
+    "${COREX_ROOT}/lib/python3/dist-packages/ixformer"/*.so \
+    "${COREX_ROOT}/lib64/python3/dist-packages/ixformer"/*.so \
+    /usr/local/lib/python3.10/dist-packages/ixformer/*.so; do
+    if [[ -f "$sopath" ]]; then
+        IX_LIBS="${IX_LIBS} ${sopath}"
+    fi
+done
+
+# Also link against libixformer*.so in corex lib dirs
+for sopath in \
+    "${COREX_ROOT}/lib64"/libixformer*.so \
+    "${COREX_ROOT}/lib64"/lib*ixformer*.so; do
+    if [[ -f "$sopath" ]]; then
+        IX_LIBS="${IX_LIBS} ${sopath}"
+    fi
+done
+
+# Add ixformer_torch_ext if present
+for sopath in \
+    "${COREX_ROOT}/lib/python3/dist-packages/ixformer"/_ixformer_torch*.so \
+    "${COREX_ROOT}/lib64/python3/dist-packages/ixformer"/_ixformer_torch*.so; do
+    if [[ -f "$sopath" ]]; then
+        IX_LIBS="${IX_LIBS} ${sopath}"
+    fi
+done
+
+if [[ -z "$IX_LIBS" ]]; then
+    echo "[ix_bridge] WARNING: No ixformer .so files found — bridge will compile but may not link all symbols" >&2
+fi
+
+# --- Locate rpath dirs ---
+RPATH_DIRS=""
+for d in \
+    "${COREX_ROOT}/lib64" \
+    "${COREX_ROOT}/lib/python3/dist-packages/ixformer" \
+    "${COREX_ROOT}/lib64/python3/dist-packages/ixformer"; do
+    if [[ -d "$d" ]]; then
+        RPATH_DIRS="${RPATH_DIRS} -Wl,-rpath,${d}"
+    fi
+done
+
+# --- Source file ---
+SRC="${CSRC_DIR}/ix_full_bridge_v2.cpp"
+if [[ ! -f "$SRC" ]]; then
+    echo "[ix_bridge] ERROR: source not found: ${SRC}" >&2
+    exit 1
+fi
+
+OUTPUT_DIR="${SCRIPT_DIR}/prebuilt"
+mkdir -p "$OUTPUT_DIR"
+OUTPUT="${OUTPUT_DIR}/ix_full_bridge_v2.so"
+
+echo "[ix_bridge] Compiling: ${SRC}"
+echo "[ix_bridge] Compiler: ${CLANGXX}"
+echo "[ix_bridge] ixformer libs: ${IX_LIBS}"
+
+$CLANGXX \
+    -shared -fPIC -O2 -std=c++17 \
+    $PYTHON_INC \
+    $TORCH_INC \
+    $TORCH_LIB \
+    -ltorch -ltorch_cpu -ltorch_python -lc10 \
+    ${IX_LIBS} \
+    ${RPATH_DIRS} \
+    -o "$OUTPUT" \
+    "$SRC"
+
+echo "[ix_bridge] ✓ Built: ${OUTPUT}"
+ls -lh "$OUTPUT"
+
+# --- Deploy if VLLM_ROOT specified ---
+if [[ -n "$VLLM_ROOT" ]] && [[ -d "$VLLM_ROOT" ]]; then
+    mkdir -p "${VLLM_ROOT}/ex_engine"
+    cp "$OUTPUT" "${VLLM_ROOT}/ex_engine/ix_full_bridge_v2.so"
+    echo "[ix_bridge] ✓ Deployed to ${VLLM_ROOT}/ex_engine/"
+fi
+
+echo "[ix_bridge] Done"

ex_engine/build_moe_bridge.sh

@@ -0,0 +1,179 @@
+#!/usr/bin/env bash
+# build_moe_bridge.sh — Compile MoE ops + bridge into ix_moe_bridge.so
+#
+# Links against:
+#   libcuinfer.so   (cuinferCustomGemm, cuinferTopK — confirmed in symbol dump)
+#   libixformer.so  (silu_and_mul, rms_norm, flash_attn, etc — confirmed)
+#
+# Real device compiler: corex clang/16, NOT nvcc
+# Reference: ex_engine/build_ix_bridge.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+VLLM_ROOT="${1:-}"
+
+echo "[moe_bridge] Building ix_moe_bridge.so"
+echo "[moe_bridge] Script dir: ${SCRIPT_DIR}"
+
+# --- Locate sources ---
+# Support both layouts:
+#   1. SCRIPT_DIR=/workspace/ex_engine → csrc/ is direct child
+#   2. SCRIPT_DIR=/workspace/qwen3_6_scripts/ex_engine_src → csrc/ is direct child
+MOE_CU=""
+BRIDGE_CPP=""
+for base in "${SCRIPT_DIR}" "${SCRIPT_DIR}/ex_engine"; do
+    [[ -f "${base}/csrc/moe_ops_impl.cu" ]] && MOE_CU="${base}/csrc/moe_ops_impl.cu"
+    [[ -f "${base}/csrc/ix_full_bridge_v2.cpp" ]] && BRIDGE_CPP="${base}/csrc/ix_full_bridge_v2.cpp"
+done
+
+if [[ -z "$MOE_CU" ]]; then
+    echo "[moe_bridge] ERROR: moe_ops_impl.cu not found under ${SCRIPT_DIR}" >&2
+    exit 1
+fi
+if [[ -z "$BRIDGE_CPP" ]]; then
+    echo "[moe_bridge] ERROR: ix_full_bridge_v2.cpp not found under ${SCRIPT_DIR}" >&2
+    exit 1
+fi
+echo "[moe_bridge] MOE_CU: ${MOE_CU}"
+echo "[moe_bridge] BRIDGE_CPP: ${BRIDGE_CPP}"
+
+# --- Locate libraries ---
+COREX_ROOT="${COREX_ROOT:-/usr/local/corex}"
+
+# Find libcuinfer.so
+CUINFER_SO=""
+for d in "${COREX_ROOT}/lib64" "${COREX_ROOT}/lib" "/usr/lib64" "/usr/lib"; do
+    if [[ -f "${d}/libcuinfer.so" ]]; then
+        CUINFER_SO="${d}/libcuinfer.so"
+        break
+    fi
+done
+
+# Find libixformer.so and ixformer Python package
+IX_LIB_DIR=""
+IX_SO_FILES=()
+for d in \
+    "${COREX_ROOT}/lib/python3/dist-packages/ixformer" \
+    "${COREX_ROOT}/lib64/python3/dist-packages/ixformer" \
+    "$(python3 -c 'import ixformer, os; print(os.path.dirname(ixformer.__file__))' 2>/dev/null || echo '')"; do
+    if [[ -d "$d" ]]; then
+        IX_LIB_DIR="$d"
+        while IFS= read -r so; do
+            IX_SO_FILES+=("$so")
+        done < <(find "$d" -name "*.so" -type f 2>/dev/null)
+        break
+    fi
+done
+
+echo "[moe_bridge] COREX_ROOT: ${COREX_ROOT}"
+echo "[moe_bridge] cuinfer: ${CUINFER_SO:-NOT FOUND}"
+echo "[moe_bridge] ixformer dir: ${IX_LIB_DIR:-NOT FOUND}"
+echo "[moe_bridge] ixformer .so count: ${#IX_SO_FILES[@]}"
+
+# --- Build via torch.utils.cpp_extension ---
+mkdir -p "${SCRIPT_DIR}/prebuilt"
+
+export SCRIPT_DIR VLLM_ROOT
+python3 << 'PYEOF'
+import os, sys, glob, shutil
+
+script_dir = os.environ.get("SCRIPT_DIR", ".")
+vllm_root = os.environ.get("VLLM_ROOT", "")
+
+# Find source files — try direct csrc/ first, then ex_engine/csrc/
+moe_cu = ""
+bridge_cpp = ""
+for base in [script_dir, os.path.join(script_dir, "ex_engine")]:
+    candidate_cu = os.path.join(base, "csrc", "moe_ops_impl.cu")
+    candidate_cpp = os.path.join(base, "csrc", "ix_full_bridge_v2.cpp")
+    if os.path.isfile(candidate_cu):
+        moe_cu = candidate_cu
+    if os.path.isfile(candidate_cpp):
+        bridge_cpp = candidate_cpp
+if not moe_cu or not bridge_cpp:
+    print(f"[moe_bridge] ERROR: sources not found under {script_dir}")
+    sys.exit(1)
+print(f"[moe_bridge] MOE_CU: {moe_cu}")
+print(f"[moe_bridge] BRIDGE_CPP: {bridge_cpp}")
+
+# Collect linker flags
+extra_ldflags = []
+rpath_dirs = set()
+
+corex_root = os.environ.get("COREX_ROOT", "/usr/local/corex")
+for search_dir in [
+    os.path.join(corex_root, "lib64"),
+    os.path.join(corex_root, "lib"),
+]:
+    if os.path.isdir(search_dir):
+        rpath_dirs.add(search_dir)
+        for so in glob.glob(os.path.join(search_dir, "libcuinfer*.so*")):
+            extra_ldflags.append(so)
+
+# ixformer .so files
+try:
+    import ixformer
+    ix_dir = os.path.dirname(ixformer.__file__)
+    rpath_dirs.add(ix_dir)
+    for so in glob.glob(os.path.join(ix_dir, "*.so")):
+        extra_ldflags.append(so)
+    for so in glob.glob(os.path.join(ix_dir, "lib*.so")):
+        if so not in extra_ldflags:
+            extra_ldflags.append(so)
+except ImportError:
+    # Search common paths
+    for d in [
+        os.path.join(corex_root, "lib", "python3", "dist-packages", "ixformer"),
+        os.path.join(corex_root, "lib64", "python3", "dist-packages", "ixformer"),
+    ]:
+        if os.path.isdir(d):
+            rpath_dirs.add(d)
+            for so in glob.glob(os.path.join(d, "*.so")):
+                extra_ldflags.append(so)
+
+for d in rpath_dirs:
+    extra_ldflags.append(f"-Wl,-rpath,{d}")
+
+print(f"[moe_bridge] Linking against {len(extra_ldflags)} items")
+for f in extra_ldflags[:10]:
+    print(f"  {f}")
+
+try:
+    from torch.utils.cpp_extension import load
+
+    mod = load(
+        name="ix_moe_bridge",
+        sources=[moe_cu, bridge_cpp],
+        extra_include_paths=[os.path.join(script_dir, "csrc")],
+        extra_cflags=["-O2", "-std=c++17"],
+        extra_cuda_cflags=["-O2", ],
+        extra_ldflags=extra_ldflags,
+        verbose=True,
+    )
+    print("[moe_bridge] ✓ Compilation successful")
+
+    # Find and copy the built .so
+    import importlib
+    spec = importlib.util.find_spec("ix_moe_bridge")
+    if spec and spec.origin:
+        dst = os.path.join(script_dir, "prebuilt", "ix_moe_bridge.so")
+        shutil.copy2(spec.origin, dst)
+        print(f"[moe_bridge] ✓ Saved to {dst}")
+
+        if vllm_root:
+            vllm_dst = os.path.join(vllm_root, "ex_engine", "ix_moe_bridge.so")
+            os.makedirs(os.path.dirname(vllm_dst), exist_ok=True)
+            shutil.copy2(spec.origin, vllm_dst)
+            print(f"[moe_bridge] ✓ Deployed to {vllm_dst}")
+    else:
+        print("[moe_bridge] ⚠ Could not locate compiled .so via importlib")
+
+except Exception as e:
+    print(f"[moe_bridge] ERROR: {e}", file=sys.stderr)
+    import traceback; traceback.print_exc()
+    sys.exit(1)
+PYEOF
+
+echo "[moe_bridge] Done"

ex_engine/build_xllm_ilu_kernels.sh

@@ -0,0 +1,127 @@
+#!/usr/bin/env bash
+# build_xllm_ilu_kernels.sh — Compile xllm upstream ILU kernel wrappers
+#
+# Source:  upstream_ref/xllm/xllm/core/kernels/ilu/*.cpp
+# Already: ex_engine/xllm_kernels/ilu/ (copied from upstream)
+# Header:  upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h
+#
+# These .cpp files are thin wrappers that call ixformer::infer C++ functions.
+# They're already proven to work on BI-V100 (xllm uses them in production).
+# We compile them into xllm_ilu_ops.so with pybind11 bindings.
+#
+# Usage:
+#   bash build_xllm_ilu_kernels.sh [VLLM_ROOT]
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+
+# Source locations — prefer ex_engine copy, fall back to upstream_ref
+ILU_DIR="${SCRIPT_DIR}/xllm_kernels/ilu"
+if [[ ! -d "$ILU_DIR" ]]; then
+    ILU_DIR="${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/ilu"
+fi
+
+if [[ ! -d "$ILU_DIR" ]]; then
+    echo "[xllm_ilu] ERROR: ILU kernel source not found" >&2
+    exit 1
+fi
+
+# Header with ixformer::infer declarations
+IXFORMER_H="${ILU_DIR}/ixformer.h"
+if [[ ! -f "$IXFORMER_H" ]]; then
+    # Copy from upstream
+    cp "${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/ilu/ixformer.h" \
+       "${ILU_DIR}/ixformer.h" 2>/dev/null || true
+    cp "${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/ilu/utils.h" \
+       "${ILU_DIR}/utils.h" 2>/dev/null || true
+fi
+
+echo "[xllm_ilu] Source dir: ${ILU_DIR}"
+echo "[xllm_ilu] Files:"
+ls -la "$ILU_DIR"/*.cpp "$ILU_DIR"/*.h 2>/dev/null || true
+
+# --- Compile via torch.utils.cpp_extension ---
+VLLM_ROOT="${1:-}"
+
+python3 << PYEOF
+import os
+import sys
+import glob
+
+# Set up paths
+ilu_dir = "${ILU_DIR}"
+script_dir = "${SCRIPT_DIR}"
+vllm_root = "${VLLM_ROOT}" if "${VLLM_ROOT}" else None
+
+# Find all .cpp files in the ILU directory
+cpp_files = sorted(glob.glob(os.path.join(ilu_dir, "*.cpp")))
+if not cpp_files:
+    print("[xllm_ilu] ERROR: No .cpp files found in", ilu_dir)
+    sys.exit(1)
+
+print(f"[xllm_ilu] Found {len(cpp_files)} source files:")
+for f in cpp_files:
+    print(f"  {os.path.basename(f)}")
+
+# Find ixformer .so files for linking
+corex_root = os.environ.get("COREX_ROOT", "/usr/local/corex")
+ix_so_files = []
+rpath_dirs = set()
+for search_dir in [
+    os.path.join(corex_root, "lib", "python3", "dist-packages", "ixformer"),
+    os.path.join(corex_root, "lib64", "python3", "dist-packages", "ixformer"),
+    os.path.join(corex_root, "lib64"),
+]:
+    if os.path.isdir(search_dir):
+        rpath_dirs.add(search_dir)
+        for so in glob.glob(os.path.join(search_dir, "*.so")):
+            ix_so_files.append(so)
+        for so in glob.glob(os.path.join(search_dir, "lib*.so")):
+            if so not in ix_so_files:
+                ix_so_files.append(so)
+
+extra_ldflags = list(ix_so_files)
+for d in rpath_dirs:
+    extra_ldflags.append(f"-Wl,-rpath,{d}")
+
+print(f"[xllm_ilu] Linking against {len(ix_so_files)} ixformer .so files")
+
+try:
+    from torch.utils.cpp_extension import load
+    mod = load(
+        name="xllm_ilu_ops",
+        sources=cpp_files,
+        extra_include_paths=[ilu_dir],
+        extra_cflags=["-O2", "-std=c++17"],
+        extra_ldflags=extra_ldflags,
+        verbose=True,
+    )
+    print("[xllm_ilu] ✓ Compilation successful")
+
+    # Save the .so
+    import torch
+    so_path = os.path.join(script_dir, "prebuilt", "xllm_ilu_ops.so")
+    os.makedirs(os.path.dirname(so_path), exist_ok=True)
+
+    # Find the compiled .so in the torch cache
+    import importlib
+    spec = importlib.util.find_spec("xllm_ilu_ops")
+    if spec and spec.origin:
+        import shutil
+        shutil.copy2(spec.origin, so_path)
+        print(f"[xllm_ilu] ✓ Saved to {so_path}")
+
+        if vllm_root:
+            dst = os.path.join(vllm_root, "ex_engine", "xllm_ilu_ops.so")
+            os.makedirs(os.path.dirname(dst), exist_ok=True)
+            shutil.copy2(spec.origin, dst)
+            print(f"[xllm_ilu] ✓ Deployed to {dst}")
+
+except Exception as e:
+    print(f"[xllm_ilu] ERROR: {e}")
+    sys.exit(1)
+PYEOF
+
+echo "[xllm_ilu] Done"

ex_engine/build_xllm_kernels.sh

@@ -0,0 +1,157 @@
+#!/usr/bin/env bash
+# build_xllm_kernels.sh — Compile xllm CUDA kernels into .so for BI-V100
+#
+# Architecture (CCCL compile pattern):
+#   CCCL: CMakePresets.json → cmake --preset cub-cpp20 → ninja → .so
+#   EX:   torch.utils.cpp_extension → clang --cuda-gpu-arch=ivcore10 → .so
+#
+# Usage:
+#   bash ex_engine/build_xllm_kernels.sh [--output-dir /path/to/output]
+#
+# Prerequisites:
+#   - BI-V100 machine with corex SDK
+#   - PyTorch with CUDA support
+#   - corex clang/16 compiler
+#
+# Outputs:
+#   xllm_fused_qknorm_rope.so — Fused QK-Norm + RoPE (saves 128 kernel launches/fwd)
+
+set -eo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+KERNELS_DIR="${SCRIPT_DIR}/xllm_kernels/cuda"
+HEADERS_DIR="${KERNELS_DIR}/headers"
+BINDINGS_DIR="${KERNELS_DIR}/bindings"
+OUTPUT_DIR="${1:-${SCRIPT_DIR}/../qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10}"
+
+mkdir -p "${OUTPUT_DIR}"
+
+echo "[build] KERNELS_DIR=${KERNELS_DIR}"
+echo "[build] HEADERS_DIR=${HEADERS_DIR}"
+echo "[build] OUTPUT_DIR=${OUTPUT_DIR}"
+
+# Common compile flags for BI-V100 (ivcore10 = SM70-class)
+CUDA_FLAGS="-O2 --cuda-gpu-arch=ivcore10"
+CXX_FLAGS="-O2 -std=c++17"
+INCLUDE_FLAGS="-I${HEADERS_DIR}"
+
+# Use torch's cpp_extension for JIT compile
+build_so() {
+    local name=$1
+    local sources=$2
+    local extra_flags="${3:-}"
+
+    echo "[build] Building ${name}.so from: ${sources}"
+
+    python3 -c "
+import os, sys
+from torch.utils.cpp_extension import load
+
+sources = '${sources}'.split()
+abs_sources = [os.path.join('${SCRIPT_DIR}', '..', s) if not os.path.isabs(s) else s for s in sources]
+abs_sources = [os.path.abspath(s) for s in abs_sources]
+
+for s in abs_sources:
+    if not os.path.exists(s):
+        print(f'ERROR: source not found: {s}', file=sys.stderr)
+        sys.exit(1)
+
+try:
+    mod = load(
+        name='${name}',
+        sources=abs_sources,
+        extra_cuda_cflags=['-O2'],
+        extra_cflags=['-O2', '-std=c++17'],
+        extra_include_paths=['${HEADERS_DIR}'],
+        build_directory='/tmp/build_${name}',
+        verbose=True,
+    )
+    # Find the compiled .so
+    import glob
+    sos = glob.glob('/tmp/build_${name}/${name}*.so')
+    if sos:
+        import shutil
+        dst = os.path.join('${OUTPUT_DIR}', '${name}.so')
+        shutil.copy2(sos[0], dst)
+        print(f'[build] SUCCESS: {dst}')
+    else:
+        print('[build] WARN: .so not found after build', file=sys.stderr)
+except Exception as e:
+    print(f'[build] FAIL ${name}: {e}', file=sys.stderr)
+    sys.exit(1)
+" || echo "[build] FAILED: ${name}"
+}
+
+# ============================================================================
+# Build targets
+# ============================================================================
+
+# 1. xllm_fused_qknorm_rope — Fused QK-Norm + RoPE
+#    Source: upstream xllm fused_qknorm_rope.cu
+#    Note: Requires corex_compat_utils.h instead of glog-dependent utils.h
+#    The .cu includes "cuda_ops_api.h" and "utils.h" — we need to make sure
+#    the include path resolves to our corex-compat headers first.
+echo ""
+echo "============================================================"
+echo "  1. xllm_fused_qknorm_rope.so"
+echo "============================================================"
+build_so "xllm_fused_qknorm_rope" \
+    "ex_engine/xllm_kernels/cuda/fused_qknorm_rope.cu ex_engine/xllm_kernels/cuda/bindings/xllm_fused_qknorm_rope_bind.cpp"
+
+# 2. xllm_norm — RMSNorm + Fused Add RMSNorm
+#    Source: upstream xllm norm.cu
+#    Hot path: called 2× per decoder layer = 72× per forward pass
+echo ""
+echo "============================================================"
+echo "  2. xllm_norm.so"
+echo "============================================================"
+build_so "xllm_norm" \
+    "ex_engine/xllm_kernels/cuda/norm.cu ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp"
+
+# 3. xllm_rope — Rotary Position Embedding
+#    Source: upstream xllm rope.cu
+#    Hot path: called 1× per attention layer = 36× per forward pass
+echo ""
+echo "============================================================"
+echo "  3. xllm_rope.so"
+echo "============================================================"
+build_so "xllm_rope" \
+    "ex_engine/xllm_kernels/cuda/rope.cu ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp"
+
+# 4. xllm_activation — SiLU-and-Mul fused activation
+#    Source: upstream xllm activation.cu
+#    Hot path: called 1× per MLP = 36× per forward pass
+echo ""
+echo "============================================================"
+echo "  4. xllm_activation.so"
+echo "============================================================"
+build_so "xllm_activation" \
+    "ex_engine/xllm_kernels/cuda/activation.cu ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp"
+
+# 5. xllm_cache — Reshape + block copy for KV cache
+#    Source: upstream xllm reshape_paged_cache.cu + block_copy.cu
+#    Hot path: called every prefill + decode step
+echo ""
+echo "============================================================"
+echo "  5. xllm_cache.so"
+echo "============================================================"
+build_so "xllm_cache" \
+    "ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu ex_engine/xllm_kernels/cuda/block_copy.cu ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp"
+
+# 6. xllm_moe — MoE topk + index + combine + fused pipeline
+#    Source: upstream xllm moe_fused_topk.cu + moe_compute_index.cu + moe_combine.cu + fused_moe.cpp
+#    THE critical .so: replaces Python for-loop over 64 experts
+echo ""
+echo "============================================================"
+echo "  6. xllm_moe.so"
+echo "============================================================"
+build_so "xllm_moe" \
+    "ex_engine/xllm_kernels/cuda/moe/moe_fused_topk.cu ex_engine/xllm_kernels/cuda/moe/moe_compute_index.cu ex_engine/xllm_kernels/cuda/moe/moe_combine.cu ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp ex_engine/xllm_kernels/cuda/bindings/xllm_moe_bind.cpp"
+
+echo ""
+echo "============================================================"
+echo "  Build complete. Output:"
+echo "============================================================"
+ls -la "${OUTPUT_DIR}"/*.so 2>/dev/null | tail -30
+echo ""
+echo "Total .so count: $(ls "${OUTPUT_DIR}"/*.so 2>/dev/null | wc -l)"

ex_engine/csrc/build_test_moe_tcu.sh

@@ -0,0 +1,160 @@
+#!/bin/bash
+# build_test_moe_tcu.sh — Build and test moe_tcu_dispatch.cpp
+set -eo pipefail
+
+echo "=== Compile moe_tcu_dispatch ==="
+python3 -c "
+import torch.utils.cpp_extension as ext
+import os, shutil, glob
+
+name = 'moe_tcu_dispatch'
+build_dir = 'ex_engine/csrc/build/tmp_' + name
+os.makedirs(build_dir, exist_ok=True)
+
+mod = ext.load(
+    name=name,
+    sources=['ex_engine/csrc/moe_tcu_dispatch.cpp'],
+    extra_cflags=['-O2', '-std=c++17'],
+    build_directory=build_dir,
+    verbose=True,
+)
+
+built = glob.glob(build_dir + '/' + name + '*.so')
+if built:
+    dst = 'ex_engine/csrc/build/' + name + '.so'
+    os.makedirs('ex_engine/csrc/build', exist_ok=True)
+    shutil.copy2(built[0], dst)
+    print(f'[build] SUCCESS: {dst}')
+"
+
+echo ""
+echo "=== Test ==="
+python3 << 'PYTEST'
+import torch
+import torch.nn.functional as F
+import sys, os, glob, time, importlib.util
+
+build_dir = 'ex_engine/csrc/build'
+so = glob.glob(f'{build_dir}/tmp_moe_tcu_dispatch/moe_tcu_dispatch*.so')
+if not so:
+    print("SKIP: .so not found")
+    sys.exit(0)
+spec = importlib.util.spec_from_file_location("moe_tcu_dispatch", so[0])
+mod = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(mod)
+print(f"Loaded: {so[0]}")
+
+# ============================================================
+# Test 1: moe_decode correctness
+# ============================================================
+print("\n--- moe_decode correctness ---")
+K, I = 128, 256
+E = 8
+top_k = 4
+hidden = torch.randn(1, K, dtype=torch.float16, device='cuda')
+w13 = torch.randn(E, 2*I, K, dtype=torch.float16, device='cuda') * 0.01
+w2 = torch.randn(E, K, I, dtype=torch.float16, device='cuda') * 0.01
+expert_ids = torch.tensor([0, 3, 5, 7], dtype=torch.int64, device='cuda')
+expert_weights = torch.tensor([0.3, 0.25, 0.25, 0.2], dtype=torch.float32, device='cuda')
+
+# C++ result
+out_cpp = mod.moe_decode(hidden, w13, w2, expert_ids, expert_weights)
+
+# Python reference
+out_py = torch.zeros_like(hidden)
+for k in range(top_k):
+    eid = expert_ids[k].item()
+    w = expert_weights[k].item()
+    gate_up = F.linear(hidden, w13[eid])
+    gate = F.silu(gate_up[:, :I])
+    up = gate_up[:, I:]
+    act = gate * up
+    expert_out = F.linear(act, w2[eid])
+    out_py += w * expert_out
+
+diff = (out_cpp.float() - out_py.float()).abs().max().item()
+print(f"  max_diff={diff:.6f} {'PASS' if diff < 1.0 else 'FAIL'}")
+
+# ============================================================
+# Test 2: moe_expert_gemm_tcu correctness
+# ============================================================
+print("\n--- moe_expert_gemm_tcu correctness ---")
+num_experts = 4
+K, N = 128, 256
+expert_counts = torch.tensor([8, 0, 16, 4], dtype=torch.int64, device='cuda')
+total = expert_counts.sum().item()
+inp = torch.randn(total, K, dtype=torch.float16, device='cuda') * 0.1
+weights = torch.randn(num_experts, N, K, dtype=torch.float16, device='cuda') * 0.1
+
+out_cpp = mod.moe_expert_gemm_tcu(inp, weights, expert_counts)
+
+# Python reference
+out_py = torch.zeros(total, N, dtype=torch.float16, device='cuda')
+off = 0
+for e in range(num_experts):
+    cnt = expert_counts[e].item()
+    if cnt == 0: continue
+    out_py[off:off+cnt] = F.linear(inp[off:off+cnt], weights[e])
+    off += cnt
+
+diff = (out_cpp.float() - out_py.float()).abs().max().item()
+print(f"  max_diff={diff:.6f} {'PASS' if diff < 0.5 else 'FAIL'}")
+
+# ============================================================
+# Test 3: Performance — Python loop vs C++ loop
+# ============================================================
+print("\n--- Performance: decode (1 token, 8 experts) ---")
+K, I = 4096, 11008
+E, top_k = 64, 8
+hidden = torch.randn(1, K, dtype=torch.float16, device='cuda')
+w13 = torch.randn(E, 2*I, K, dtype=torch.float16, device='cuda') * 0.001
+w2 = torch.randn(E, K, I, dtype=torch.float16, device='cuda') * 0.001
+expert_ids = torch.tensor([0,5,10,20,30,40,50,60], dtype=torch.int64, device='cuda')
+expert_weights = torch.ones(top_k, dtype=torch.float32, device='cuda') / top_k
+
+# Warmup
+for _ in range(3):
+    mod.moe_decode(hidden, w13, w2, expert_ids, expert_weights)
+torch.cuda.synchronize()
+
+# C++ loop
+t0 = time.time()
+for _ in range(100):
+    mod.moe_decode(hidden, w13, w2, expert_ids, expert_weights)
+torch.cuda.synchronize()
+ms_cpp = (time.time() - t0) / 100 * 1000
+
+# Python loop
+for _ in range(3):
+    out_py = torch.zeros_like(hidden)
+    for k in range(top_k):
+        eid = expert_ids[k].item()
+        w = expert_weights[k].item()
+        gate_up = F.linear(hidden, w13[eid])
+        gate = F.silu(gate_up[:, :I])
+        up = gate_up[:, I:]
+        act = gate * up
+        out_py += w * F.linear(act, w2[eid])
+torch.cuda.synchronize()
+
+t0 = time.time()
+for _ in range(100):
+    out_py = torch.zeros_like(hidden)
+    for k in range(top_k):
+        eid = expert_ids[k].item()
+        w = expert_weights[k].item()
+        gate_up = F.linear(hidden, w13[eid])
+        gate = F.silu(gate_up[:, :I])
+        up = gate_up[:, I:]
+        act = gate * up
+        out_py += w * F.linear(act, w2[eid])
+torch.cuda.synchronize()
+ms_py = (time.time() - t0) / 100 * 1000
+
+print(f"  C++ loop:    {ms_cpp:.2f} ms")
+print(f"  Python loop: {ms_py:.2f} ms")
+print(f"  Speedup:     {ms_py/ms_cpp:.2f}x")
+print(f"  Saved:       {ms_py-ms_cpp:.2f} ms per forward")
+
+print("\n=== DONE ===")
+PYTEST

ex_engine/csrc/cuinfer_gemm_wrapper.cu

@@ -0,0 +1,161 @@
+// cuinfer_gemm_wrapper.cu — Wrapper around cuinferCustomGemm
+//
+// ixformer::functions::cuinfer_gemm exists in libixformer.so but
+// takes ixformer::Tensor (not torch::Tensor). We need a torch-compatible
+// wrapper that calls the C API directly.
+//
+// Symbol dump shows cuinferCustomGemm in libcuinfer.so with signature:
+//   cuinferCustomGemm(handle, stream, ptrMode, transa, transb,
+//                     m, n, k, alpha, A, Atype, lda, strideA,
+//                     B, Btype, ldb, strideB, beta,
+//                     C, Ctype, ldc, strideC, batchCount,
+//                     computeType, scaleType, customHostPtr, customDevicePtr, customOption)
+//
+// Reference:
+//   cat_files/ixinfer.h — cuinferCustomGemm signature
+//   libixformer.so — ixformer::functions::cuinfer_gemm (confirmed in symbol dump)
+
+#include <torch/extension.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_fp16.h>
+#include "cuinfer_handle.h"
+
+// cuinferCustomGemm is already declared in cuinfer_handle.h extern "C" block
+// We add the full signature here
+extern "C" {
+int cuinferCustomGemm(
+    cuinferHandle_t handle, cudaStream_t stream,
+    int ptrMode, int transa, int transb,
+    int m, int n, int k,
+    const void* alpha,
+    const void* A, int Atype, int lda, long long int strideA,
+    const void* B, int Btype, int ldb, long long int strideB,
+    const void* beta,
+    void* C, int Ctype, int ldc, long long int strideC,
+    int batchCount, int computeType, int scaleType,
+    const void* customHostPtr, const void* customDevicePtr, int customOption);
+}
+
+// CUDA_R_16F = 2, CUDA_R_32F = 0 (from cudaDataType_t)
+static constexpr int kFP16 = 2;
+static constexpr int kFP32 = 0;
+
+
+// ============================================================================
+// cuinfer_gemm: C = alpha * A @ B + beta * C
+//
+// A: (M, K) row-major fp16
+// B: (K, N) row-major fp16  (or (N, K) if transb)
+// C: (M, N) row-major fp16
+// ============================================================================
+torch::Tensor cuinfer_gemm(
+    torch::Tensor A,       // (M, K)
+    torch::Tensor B,       // (K, N) or (N, K) if trans_b
+    bool trans_b)
+{
+    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "inputs must be CUDA");
+    TORCH_CHECK(A.scalar_type() == torch::kHalf, "A must be fp16");
+    TORCH_CHECK(B.scalar_type() == torch::kHalf, "B must be fp16");
+
+    int M = A.size(0);
+    int K = A.size(1);
+    int N = trans_b ? B.size(0) : B.size(1);
+
+    if (!trans_b) {
+        TORCH_CHECK(B.size(0) == K, "B rows must equal K");
+    } else {
+        TORCH_CHECK(B.size(1) == K, "B cols must equal K when transposed");
+    }
+
+    auto C = torch::zeros({M, N}, A.options());
+    auto stream = c10::cuda::getCurrentCUDAStream().stream();
+    auto handle = CuinferHandle::get(stream);
+
+    if (!handle) {
+        // Fallback to torch::mm
+        if (trans_b) {
+            return torch::mm(A.to(torch::kFloat32), B.t().to(torch::kFloat32)).to(torch::kHalf);
+        }
+        return torch::mm(A.to(torch::kFloat32), B.to(torch::kFloat32)).to(torch::kHalf);
+    }
+
+    float alpha = 1.0f, beta = 0.0f;
+    int transa = 0;  // N = no transpose
+    int transb_flag = trans_b ? 1 : 0;
+
+    int lda = K;
+    int ldb = trans_b ? K : N;
+    int ldc = N;
+
+    int status = cuinferCustomGemm(
+        handle, stream,
+        0,                           // CUINFER_POINTER_MODE_HOST
+        transa, transb_flag,
+        M, N, K,
+        &alpha,
+        A.data_ptr(), kFP16, lda, 0,
+        B.data_ptr(), kFP16, ldb, 0,
+        &beta,
+        C.data_ptr(), kFP16, ldc, 0,
+        1,                           // batchCount
+        kFP32, kFP32,                // computeType, scaleType
+        nullptr, nullptr, 0);
+
+    TORCH_CHECK(status == 0, "cuinferCustomGemm failed with status ", status);
+    return C;
+}
+
+
+// ============================================================================
+// cuinfer_gemm_batched: batched version
+// A: (batch, M, K), B: (batch, K, N) or (batch, N, K)
+// ============================================================================
+torch::Tensor cuinfer_gemm_batched(
+    torch::Tensor A,
+    torch::Tensor B,
+    bool trans_b)
+{
+    TORCH_CHECK(A.dim() == 3 && B.dim() == 3, "inputs must be 3D");
+
+    int batch = A.size(0);
+    int M = A.size(1);
+    int K = A.size(2);
+    int N = trans_b ? B.size(1) : B.size(2);
+
+    auto C = torch::zeros({batch, M, N}, A.options());
+    auto stream = c10::cuda::getCurrentCUDAStream().stream();
+    auto handle = CuinferHandle::get(stream);
+
+    float alpha = 1.0f, beta = 0.0f;
+    int lda = K, ldb = trans_b ? K : N, ldc = N;
+    long long strideA = (long long)M * K;
+    long long strideB = trans_b ? (long long)N * K : (long long)K * N;
+    long long strideC = (long long)M * N;
+
+    int status = cuinferCustomGemm(
+        handle, stream,
+        0,
+        0, trans_b ? 1 : 0,
+        M, N, K,
+        &alpha,
+        A.data_ptr(), kFP16, lda, strideA,
+        B.data_ptr(), kFP16, ldb, strideB,
+        &beta,
+        C.data_ptr(), kFP16, ldc, strideC,
+        batch,
+        kFP32, kFP32,
+        nullptr, nullptr, 0);
+
+    TORCH_CHECK(status == 0, "cuinferCustomGemm batched failed: ", status);
+    return C;
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("cuinfer_gemm", &cuinfer_gemm,
+          "GEMM via cuinferCustomGemm (fp16, Cu10)",
+          py::arg("A"), py::arg("B"), py::arg("trans_b") = false);
+    m.def("cuinfer_gemm_batched", &cuinfer_gemm_batched,
+          "Batched GEMM via cuinferCustomGemm",
+          py::arg("A"), py::arg("B"), py::arg("trans_b") = false);
+}

ex_engine/csrc/cuinfer_handle.h

@@ -0,0 +1,65 @@
+// cuinfer_handle.h — Singleton handle manager for libcuinfer.so
+//
+// cuinferCreate/Destroy is expensive. This provides a thread-safe
+// singleton that creates once and reuses.
+//
+// Usage:
+//   #include "cuinfer_handle.h"
+//   cuinferHandle_t h = CuinferHandle::get(stream);
+//
+// Reference: ixformer::Context::default_cuinfer_handle (in libixformer.so)
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include <mutex>
+#include <cstdio>
+
+// Forward-declare cuinfer C API
+extern "C" {
+
+typedef struct cuinferContext* cuinferHandle_t;
+
+typedef enum {
+    CUINFER_STATUS_SUCCESS_H = 0,
+} cuinferStatus_h_t;
+
+int cuinferCreate(cuinferHandle_t* handle);
+int cuinferDestroy(cuinferHandle_t handle);
+int cuinferSetStream(cuinferHandle_t handle, cudaStream_t stream);
+
+}  // extern "C"
+
+
+class CuinferHandle {
+public:
+    static cuinferHandle_t get(cudaStream_t stream = nullptr) {
+        static CuinferHandle instance;
+        if (stream && stream != instance.last_stream_) {
+            cuinferSetStream(instance.handle_, stream);
+            instance.last_stream_ = stream;
+        }
+        return instance.handle_;
+    }
+
+private:
+    cuinferHandle_t handle_ = nullptr;
+    cudaStream_t last_stream_ = nullptr;
+
+    CuinferHandle() {
+        int status = cuinferCreate(&handle_);
+        if (status != 0) {
+            fprintf(stderr, "[cuinfer_handle] WARNING: cuinferCreate failed (%d)\n", status);
+            handle_ = nullptr;
+        }
+    }
+
+    ~CuinferHandle() {
+        if (handle_) {
+            cuinferDestroy(handle_);
+        }
+    }
+
+    CuinferHandle(const CuinferHandle&) = delete;
+    CuinferHandle& operator=(const CuinferHandle&) = delete;
+};

ex_engine/csrc/cuinfer_types.h

@@ -0,0 +1,175 @@
+// cuinfer_types.h — C API types from libcuinfer.so
+//
+// Extracted from: cat_files/ixinfer.h (165952 bytes, from real device)
+// Only the types/enums needed by our GEMM and MoE code.
+//
+// This header replaces the scattered extern "C" blocks across
+// moe_ops_impl.cu, cuinfer_gemm_wrapper.cu, gemm_grouped.cu.
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// --- Handle ---
+struct cuinferContext;
+typedef struct cuinferContext* cuinferHandle_t;
+
+// --- Status ---
+typedef enum {
+    CUINFER_STATUS_SUCCESS = 0,
+    CUINFER_STATUS_NOT_INITIALIZED = 1,
+    CUINFER_STATUS_ALLOC_FAILED = 2,
+    CUINFER_STATUS_BAD_PARAM = 3,
+    CUINFER_STATUS_INTERNAL_ERROR = 4,
+    CUINFER_STATUS_INVALID_VALUE = 5,
+    CUINFER_STATUS_ARCH_MISMATCH = 6,
+    CUINFER_STATUS_EXECUTION_FAILED = 8,
+    CUINFER_STATUS_NOT_SUPPORTED = 9,
+} cuinferStatus_t;
+
+// --- Data types ---
+typedef enum {
+    CUINFER_DATA_FLOAT = 0,
+    CUINFER_DATA_DOUBLE = 1,
+    CUINFER_DATA_HALF = 2,
+    CUINFER_DATA_INT8 = 3,
+    CUINFER_DATA_INT32 = 4,
+    CUINFER_DATA_INT8x4 = 5,
+    CUINFER_DATA_UINT8 = 6,
+    CUINFER_DATA_UINT8x4 = 7,
+    CUINFER_DATA_INT16 = 8,
+    CUINFER_DATA_BFLOAT16 = 9,
+} cuinferDataType_t;
+
+// --- Operations ---
+typedef enum {
+    CUINFER_OP_N = 0,   // no transpose
+    CUINFER_OP_T = 1,   // transpose
+    CUINFER_OP_C = 2,   // conjugate transpose
+} cuinferOperation_t;
+
+// --- Pointer mode ---
+typedef enum {
+    CUINFER_POINTER_MODE_HOST = 0,
+    CUINFER_POINTER_MODE_DEVICE = 1,
+} cuinferPointerMode_t;
+
+// --- GEMM custom option ---
+typedef enum {
+    CUINFER_GEMM_DEFAULT = 0,
+} cuinferGEMMCustomOption_t;
+
+// --- Reduce ops ---
+typedef enum {
+    CUINFER_REDUCE_TENSOR_ADD = 0,
+    CUINFER_REDUCE_TENSOR_MUL = 1,
+    CUINFER_REDUCE_TENSOR_MIN = 2,
+    CUINFER_REDUCE_TENSOR_MAX = 3,
+} cuinferReduceTensorOp_t;
+
+// --- Softmax ---
+typedef enum {
+    CUINFER_SOFTMAX_FAST = 0,
+    CUINFER_SOFTMAX_ACCURATE = 1,
+    CUINFER_SOFTMAX_LOG = 2,
+} cuinferSoftmaxAlgorithm_t;
+
+typedef enum {
+    CUINFER_SOFTMAX_MODE_INSTANCE = 0,
+    CUINFER_SOFTMAX_MODE_CHANNEL = 1,
+} cuinferSoftmaxMode_t;
+
+
+// ============================================================================
+// Function declarations (confirmed in libcuinfer.so symbol dump)
+// ============================================================================
+
+cuinferStatus_t cuinferCreate(cuinferHandle_t* handle);
+cuinferStatus_t cuinferDestroy(cuinferHandle_t handle);
+cuinferStatus_t cuinferSetStream(cuinferHandle_t handle, cudaStream_t stream);
+cuinferStatus_t cuinferGetStream(cuinferHandle_t handle, cudaStream_t* stream);
+size_t cuinferGetVersion(void);
+const char* cuinferGetErrorString(cuinferStatus_t status);
+
+// GEMM
+cuinferStatus_t cuinferCustomGemm(
+    cuinferHandle_t handle, cudaStream_t stream,
+    cuinferPointerMode_t ptrMode,
+    cuinferOperation_t transa, cuinferOperation_t transb,
+    int m, int n, int k,
+    const void* alpha,
+    const void* A, cudaDataType_t Atype, int lda, long long int strideA,
+    const void* B, cudaDataType_t Btype, int ldb, long long int strideB,
+    const void* beta,
+    void* C, cudaDataType_t Ctype, int ldc, long long int strideC,
+    int batchCount,
+    cudaDataType_t computeType, cudaDataType_t scaleType,
+    const void* customHostPtr, const void* customDevicePtr,
+    cuinferGEMMCustomOption_t customOption);
+
+cuinferStatus_t cuinferCustomGemmEx(
+    cuinferHandle_t handle, cudaStream_t stream,
+    cuinferPointerMode_t ptrMode,
+    cuinferOperation_t transa, cuinferOperation_t transb,
+    int m, int n, int k,
+    const void* alpha,
+    const void* A, cudaDataType_t Atype, int lda, long long int strideA,
+    const void* B, cudaDataType_t Btype, int ldb, long long int strideB,
+    const void* beta,
+    void* C, cudaDataType_t Ctype, int ldc, long long int strideC,
+    int batchCount,
+    cudaDataType_t computeType, cudaDataType_t scaleType,
+    const void* customHostPtr, const void* customDevicePtr,
+    cuinferGEMMCustomOption_t customOption,
+    const void* workspace);
+
+// TopK
+cuinferStatus_t cuinferTopK(
+    cuinferHandle_t handle,
+    const void* input, int n, int m, int top_k,
+    int sort_dim, bool largest, bool sorted,
+    void* out_value, int* out_indice,
+    cuinferDataType_t datatype, void* workspace);
+
+cuinferStatus_t cuinferGetTopKWorkspace(
+    cuinferHandle_t handle,
+    int n, int m, int top_k,
+    cuinferDataType_t datatype, size_t* workspace_size);
+
+cuinferStatus_t cuinferTopKBatch(
+    cuinferHandle_t handle,
+    const void* input, int top_k, int batch, int n, int m, int k,
+    bool largest, bool sorted, int sort_dim,
+    void* output, int* indice,
+    cuinferDataType_t datatype, void* workspace);
+
+// Softmax
+cuinferStatus_t cuinferSoftmaxForward(
+    cuinferHandle_t handle,
+    cuinferSoftmaxAlgorithm_t algo,
+    cuinferSoftmaxMode_t mode,
+    const void* alpha,
+    const void* xDesc, const void* x,
+    const void* beta,
+    const void* yDesc, void* y);
+
+// Reduce
+cuinferStatus_t cuinferReduce(
+    cuinferHandle_t handle,
+    const void* in, void* out,
+    cuinferDataType_t in_type,
+    cuinferDataType_t acc_type,
+    cuinferDataType_t out_type,
+    cuinferReduceTensorOp_t reduce_op,
+    int n_dims, const int* dims,
+    int n_reduce_dims, const int* reduce_dim_index,
+    void* workspace);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif

ex_engine/csrc/gemm_grouped.cu

@@ -0,0 +1,188 @@
+// gemm_grouped.cu — Per-expert GEMM using CUTLASS Cu10 TensorOp
+//
+// Source lineage:
+//   cat_files/batched_gemm.cu           — cutlass sample from real device
+//   cat_files/default_gemm_configuration.h — Cu10 half/half/float config
+//   ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu — existing impl
+//   ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp — moe_expert_gemm pattern
+//
+// This file provides:
+//   1. cutlass_expert_gemm()   — one cutlass GEMM per expert (Cu10 TensorOp)
+//   2. cuinfer_expert_gemm()   — one cuinferCustomGemm per expert (fallback)
+//   3. moe_group_gemm()        — unified entry: try cutlass, fall back to cuinfer
+//
+// All use RowMajor, FP16 data, FP32 accumulation.
+// Weight layout: [num_experts, N, K] (TN format = transB in GEMM sense)
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/gemm/device/gemm_batched.h"
+
+// ============================================================================
+// Cu10 TensorOp GEMM type — from default_gemm_configuration.h
+// ThreadblockShape<128,128,32>, WarpShape<32,32,32>, Instruction<16,16,16>
+// ============================================================================
+using GemmCu10 = cutlass::gemm::device::GemmBatched<
+    cutlass::half_t,                    // ElementA
+    cutlass::layout::RowMajor,          // LayoutA
+    cutlass::half_t,                    // ElementB
+    cutlass::layout::RowMajor,          // LayoutB
+    cutlass::half_t,                    // ElementC
+    cutlass::layout::RowMajor,          // LayoutC
+    float,                              // ElementAccumulator
+    cutlass::arch::OpClassTensorOp,     // use TCU
+    cutlass::arch::Cu10                 // BI-V100
+>;
+
+
+// ============================================================================
+// cutlass_expert_gemm: per-expert GEMM using CUTLASS
+//
+// For each expert e with M_e tokens:
+//   C[offset:offset+M_e, :N] = A[offset:offset+M_e, :K] @ B[e, :N, :K]^T
+//
+// B is stored as [num_experts, N, K] (RowMajor), we need A×B^T.
+// Cutlass RowMajor × RowMajor computes C = A × B, so we transpose:
+//   C(M,N) = A(M,K) × B^T(K,N) = A(M,K) × B_orig(N,K)^T
+//
+// In row-major: A lda=K, B lda=K (it's NxK stored row-major), C ldc=N
+// We use Cutlass's NN mode on (A, B^T) which is implemented as:
+//   Cutlass RowMajor NN: C[i,j] = sum_k A[i,k] * B[k,j]
+//   But B is (N,K) not (K,N), so we pass B as ColumnMajor or handle via stride.
+//
+// Simpler: A is (M,K) RowMajor, we want output (M,N).
+// B_expert is (N,K) RowMajor = same as (K,N) ColumnMajor.
+// So: A(M,K) RowMajor × B(K,N) ColumnMajor → C(M,N) RowMajor
+// This is exactly GEMM with transB.
+// ============================================================================
+
+using GemmCu10_TN = cutlass::gemm::device::GemmBatched<
+    cutlass::half_t,                    // ElementA
+    cutlass::layout::RowMajor,          // LayoutA — A is (M,K) row-major
+    cutlass::half_t,                    // ElementB
+    cutlass::layout::ColumnMajor,       // LayoutB — B is (N,K) stored row = (K,N) col
+    cutlass::half_t,                    // ElementC
+    cutlass::layout::RowMajor,          // LayoutC
+    float,                              // ElementAccumulator
+    cutlass::arch::OpClassTensorOp,     // TCU
+    cutlass::arch::Cu10                 // BI-V100
+>;
+
+
+int cutlass_expert_gemm(
+    int num_experts,
+    const int* expert_counts,     // host array [num_experts]
+    const int* expert_offsets,    // host array [num_experts], exclusive prefix sum
+    int N, int K,
+    const __half* input,          // (total_tokens, K) row-major
+    const __half* weights,        // (num_experts, N, K) row-major — TN format
+    __half* output,               // (total_tokens, N) row-major
+    cudaStream_t stream)
+{
+    GemmCu10_TN gemm_op;
+    float alpha = 1.0f, beta = 0.0f;
+    int failures = 0;
+
+    for (int e = 0; e < num_experts; e++) {
+        int M_e = expert_counts[e];
+        if (M_e <= 0) continue;
+
+        int off = expert_offsets[e];
+        auto A = reinterpret_cast<cutlass::half_t const*>(input + (long long)off * K);
+        auto B = reinterpret_cast<cutlass::half_t const*>(weights + (long long)e * N * K);
+        auto C = reinterpret_cast<cutlass::half_t*>(output + (long long)off * N);
+
+        // A: (M_e, K) RowMajor, lda = K
+        // B: (N, K) RowMajor → (K, N) ColumnMajor, ldb = N (col-major stride)
+        // C: (M_e, N) RowMajor, ldc = N
+        cutlass::Status status = gemm_op({
+            {M_e, N, K},
+            {A, K},       // A, lda
+            0,             // strideA (not batched)
+            {B, K},        // B in col-major view: (N,K) row = (K,N) col, ldb = K
+            0,             // strideB
+            {C, N},        // C, ldc
+            0,             // strideC
+            {C, N},        // D = C
+            0,
+            {alpha, beta},
+            1              // batch_count = 1 (we loop over experts)
+        });
+
+        if (status != cutlass::Status::kSuccess) {
+            failures++;
+        }
+    }
+    return failures;
+}
+
+
+// ============================================================================
+// cuinfer fallback — forward-declare cuinferCustomGemm
+// ============================================================================
+extern "C" {
+typedef struct cuinferContext* cuinferHandle_t;
+typedef enum { CUINFER_STATUS_SUCCESS_GG = 0 } cuinferStatus_gg_t;
+cuinferHandle_t cuinferCreate_handle();
+
+int cuinferCustomGemm(
+    cuinferHandle_t handle, cudaStream_t stream,
+    int ptrMode, int transa, int transb,
+    int m, int n, int k,
+    const void* alpha,
+    const void* A, int Atype, int lda, long long int strideA,
+    const void* B, int Btype, int ldb, long long int strideB,
+    const void* beta,
+    void* C, int Ctype, int ldc, long long int strideC,
+    int batchCount, int computeType, int scaleType,
+    const void* customHostPtr, const void* customDevicePtr, int customOption);
+}
+
+
+int cuinfer_expert_gemm(
+    int num_experts,
+    const int* expert_counts,
+    const int* expert_offsets,
+    int N, int K,
+    const __half* input,
+    const __half* weights,
+    __half* output,
+    cudaStream_t stream,
+    cuinferHandle_t handle)
+{
+    float alpha = 1.0f, beta = 0.0f;
+    int failures = 0;
+
+    for (int e = 0; e < num_experts; e++) {
+        int M_e = expert_counts[e];
+        if (M_e <= 0) continue;
+
+        int off = expert_offsets[e];
+        const void* A = input + (long long)off * K;
+        const void* B = weights + (long long)e * N * K;
+        void* C = output + (long long)off * N;
+
+        // cuinferCustomGemm: transa=0 (N), transb=1 (T)
+        // CUDA_R_16F = 2
+        int status = cuinferCustomGemm(
+            handle, stream,
+            0,     // CUINFER_POINTER_MODE_HOST
+            0, 1,  // transa=N, transb=T
+            M_e, N, K,
+            &alpha,
+            A, 2, K, 0,   // A: fp16, lda=K
+            B, 2, K, 0,   // B: fp16, ldb=K (row-major N×K, transposed)
+            &beta,
+            C, 2, N, 0,   // C: fp16, ldc=N
+            1,             // batchCount=1
+            0, 0,          // computeType=fp32, scaleType=fp32
+            nullptr, nullptr, 0);
+
+        if (status != 0) failures++;
+    }
+    return failures;
+}

ex_engine/csrc/gemm_grouped_bind.cpp

@@ -0,0 +1,182 @@
+// gemm_grouped_bind.cpp — Python bindings for grouped GEMM
+//
+// Source lineage:
+//   ex_engine/xllm_kernels/cuda/bindings/hgemm_bind.cpp — moe_expert_gemm pattern
+//   ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp — batched pattern
+//
+// Exports:
+//   moe_group_gemm(input, weights, expert_counts) → output
+//   moe_group_gemm_cutlass(input, weights, expert_counts) → output
+//   moe_decode_cutlass(hidden, w13, w2, topk_weights) → output
+
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <c10/cuda/CUDAStream.h>
+#include <vector>
+
+// From gemm_grouped.cu
+int cutlass_expert_gemm(
+    int num_experts,
+    const int* expert_counts, const int* expert_offsets,
+    int N, int K,
+    const __half* input, const __half* weights, __half* output,
+    cudaStream_t stream);
+
+
+// ============================================================================
+// moe_group_gemm: per-expert GEMM using CUTLASS Cu10 TensorOp
+//
+// input:         (total_tokens, K) fp16
+// weights:       (num_experts, N, K) fp16, TN layout
+// expert_counts: (num_experts,) int32
+// Returns:       (total_tokens, N) fp16
+// ============================================================================
+torch::Tensor moe_group_gemm(
+    torch::Tensor input,
+    torch::Tensor weights,
+    torch::Tensor expert_counts)
+{
+    TORCH_CHECK(input.is_cuda() && weights.is_cuda(), "inputs must be CUDA");
+    TORCH_CHECK(input.scalar_type() == torch::kHalf, "input must be fp16");
+    TORCH_CHECK(weights.scalar_type() == torch::kHalf, "weights must be fp16");
+
+    int total_tokens = input.size(0);
+    int K = input.size(1);
+    int num_experts = weights.size(0);
+    int N = weights.size(1);
+    TORCH_CHECK(weights.size(2) == K, "weights K dim must match input K");
+
+    auto output = torch::zeros({total_tokens, N}, input.options());
+
+    // Build host arrays
+    auto counts_cpu = expert_counts.to(torch::kCPU).to(torch::kInt32).contiguous();
+    int32_t* c = counts_cpu.data_ptr<int32_t>();
+    std::vector<int> counts(num_experts), offsets(num_experts);
+    int cumsum = 0;
+    for (int i = 0; i < num_experts; i++) {
+        counts[i] = c[i];
+        offsets[i] = cumsum;
+        cumsum += c[i];
+    }
+
+    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();
+
+    int fails = cutlass_expert_gemm(
+        num_experts, counts.data(), offsets.data(),
+        N, K,
+        reinterpret_cast<const __half*>(input.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(weights.data_ptr<at::Half>()),
+        reinterpret_cast<__half*>(output.data_ptr<at::Half>()),
+        stream);
+
+    if (fails > 0) {
+        // Fallback to PyTorch F.linear per expert
+        auto input_a = input.to(torch::kFloat32);
+        auto output_f = torch::zeros({total_tokens, N},
+            input.options().dtype(torch::kFloat32));
+        for (int e = 0; e < num_experts; e++) {
+            if (counts[e] <= 0) continue;
+            int off = offsets[e];
+            auto x = input_a.narrow(0, off, counts[e]);
+            auto w = weights[e].to(torch::kFloat32);  // (N, K)
+            output_f.narrow(0, off, counts[e]) = torch::mm(x, w.t());
+        }
+        output = output_f.to(torch::kHalf);
+    }
+
+    return output;
+}
+
+
+// ============================================================================
+// moe_decode_cutlass: fused MoE decode for single-token (batch=1)
+//
+// Uses CUTLASS batched GEMM for the topk experts simultaneously.
+//
+// hidden:       (1, H) fp16
+// w13_sel:      (topk, 2*I, H) fp16 — already-gathered expert weights
+// w2_sel:       (topk, H, I) fp16
+// topk_weights: (topk,) float32
+// Returns:      (1, H) fp16
+// ============================================================================
+
+// From corex_batched_gemm_kernel.cu
+cudaError_t cutlass_batched_hgemm(
+    int m, int n, int k,
+    __half const *A, int lda, long long int batch_stride_A,
+    __half const *B, int ldb, long long int batch_stride_B,
+    __half *C, int ldc, long long int batch_stride_C,
+    int batch_count);
+
+
+torch::Tensor moe_decode_cutlass(
+    torch::Tensor hidden,         // (1, H)
+    torch::Tensor w13_sel,        // (topk, 2*I, H)
+    torch::Tensor w2_sel,         // (topk, H, I)
+    torch::Tensor topk_weights)   // (topk,)
+{
+    int topk = w13_sel.size(0);
+    int two_I = w13_sel.size(1);
+    int H = w13_sel.size(2);
+    int I = two_I / 2;
+
+    // x: (1,H) → expand to (topk, 1, H)
+    auto x = hidden.expand({topk, 1, H}).contiguous();
+
+    // w13^T: (topk, 2I, H) → transpose → (topk, H, 2I)
+    auto w13_t = w13_sel.transpose(1, 2).contiguous();
+
+    // Step 1: gate_up = x @ w13^T → (topk, 1, 2I)
+    auto gate_up_3d = torch::empty({topk, 1, two_I}, x.options());
+    auto status1 = cutlass_batched_hgemm(
+        1, two_I, H,
+        reinterpret_cast<const __half*>(x.data_ptr<at::Half>()),
+        H, H,
+        reinterpret_cast<const __half*>(w13_t.data_ptr<at::Half>()),
+        two_I, H * two_I,
+        reinterpret_cast<__half*>(gate_up_3d.data_ptr<at::Half>()),
+        two_I, two_I,
+        topk);
+    TORCH_CHECK(status1 == cudaSuccess, "batched GEMM 1 failed");
+
+    auto gate_up = gate_up_3d.squeeze(1);  // (topk, 2I)
+
+    // Step 2: SiLU activation
+    auto chunks = gate_up.chunk(2, 1);
+    auto act = torch::silu(chunks[0]) * chunks[1];  // (topk, I)
+    act = act.unsqueeze(1).contiguous();  // (topk, 1, I)
+
+    // w2^T: (topk, H, I) → transpose → (topk, I, H)
+    auto w2_t = w2_sel.transpose(1, 2).contiguous();
+
+    // Step 3: down = act @ w2^T → (topk, 1, H)
+    auto down_3d = torch::empty({topk, 1, H}, x.options());
+    auto status2 = cutlass_batched_hgemm(
+        1, H, I,
+        reinterpret_cast<const __half*>(act.data_ptr<at::Half>()),
+        I, I,
+        reinterpret_cast<const __half*>(w2_t.data_ptr<at::Half>()),
+        H, I * H,
+        reinterpret_cast<__half*>(down_3d.data_ptr<at::Half>()),
+        H, H,
+        topk);
+    TORCH_CHECK(status2 == cudaSuccess, "batched GEMM 2 failed");
+
+    auto down = down_3d.squeeze(1);  // (topk, H)
+
+    // Step 4: weighted sum
+    auto out = (down * topk_weights.unsqueeze(1).to(down.dtype())).sum(0, true);
+    return out;
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("moe_group_gemm", &moe_group_gemm,
+          "Per-expert GEMM via CUTLASS Cu10 TensorOp",
+          py::arg("input"), py::arg("weights"), py::arg("expert_counts"));
+    m.def("moe_decode_cutlass", &moe_decode_cutlass,
+          "Fused MoE decode via CUTLASS batched GEMM",
+          py::arg("hidden"), py::arg("w13_sel"),
+          py::arg("w2_sel"), py::arg("topk_weights"));
+}

ex_engine/csrc/ix_attn_bridge.cpp

@@ -0,0 +1,221 @@
+// ix_attn_bridge.cpp — Bridge to ixformer::infer attention + linear functions
+//
+// Exposes functions from ixformer.h that are NOT available via ixformer.functions:
+//   1. ixinfer_flash_attn_unpad_with_block_tables — fused prefill attention
+//   2. xllm_paged_attention — fused paged decode attention
+//   3. ixformer_linear — fused linear (matmul + optional activation)
+//   4. ixformer_linear_ex — simple fused linear
+//   5. residual_rms_norm — fused residual + RMS norm (NOT in ixformer_torch_ext)
+//
+// Source: xllm/xllm/core/kernels/ilu/ixformer.h
+// Usage: xllm/xllm/core/kernels/ilu/attention.cpp
+//        xllm/xllm/core/layers/ilu/attention.cpp
+
+#include <torch/extension.h>
+#include <optional>
+
+namespace ixformer {
+namespace infer {
+
+// Prefill: flash attention with block tables (variable-length batched)
+torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    torch::Tensor& out,
+    torch::Tensor& block_tables,
+    torch::Tensor& cu_seq_q,
+    torch::Tensor& cu_seq_k,
+    int64_t max_seq_q,
+    int64_t max_seq_k,
+    bool is_causal,
+    int64_t window_left,
+    int64_t window_right,
+    double scale,
+    double softcap,
+    bool sqrt_alibi,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::optional<torch::Tensor>& sinks,
+    std::optional<torch::Tensor>& lse);
+
+// Decode: paged attention (single-step cached KV)
+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);
+
+// Fused linear: matmul + optional activation
+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);
+
+// Simple linear
+torch::Tensor ixformer_linear_ex(
+    torch::Tensor& input,
+    torch::Tensor& weight,
+    const c10::optional<torch::Tensor>& bias,
+    const c10::optional<torch::Tensor>& out);
+
+// Fused residual + RMS norm (not in ixformer_torch_ext, only in ixformer::infer)
+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);
+
+}  // namespace infer
+}  // namespace ixformer
+
+
+// ============================================================================
+// Python-facing wrappers
+// Port from: xllm/xllm/core/kernels/ilu/attention.cpp
+// ============================================================================
+
+// Prefill attention via flash_attn_unpad_with_block_tables
+torch::Tensor ix_prefill_attention(
+    torch::Tensor query,           // (total_q_tokens, num_heads, head_dim)
+    torch::Tensor key_cache,       // (num_blocks, num_heads, block_size, head_dim)
+    torch::Tensor value_cache,     // (num_blocks, num_heads, block_size, head_dim)
+    torch::Tensor output,          // (total_q_tokens, num_heads, head_dim)
+    torch::Tensor block_tables,    // (batch, max_blocks)
+    torch::Tensor cu_seq_q,        // (batch+1,)
+    torch::Tensor cu_seq_k,        // (batch+1,)
+    int64_t max_query_len,
+    int64_t max_seq_len,
+    double scale,
+    bool is_causal,
+    int64_t window_left,
+    int64_t window_right) {
+
+    std::optional<torch::Tensor> lse;
+
+    return ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables(
+        query, key_cache, value_cache, output, block_tables,
+        cu_seq_q, cu_seq_k,
+        max_query_len, max_seq_len,
+        is_causal,
+        window_left, window_right,
+        scale,
+        /*softcap=*/0.0,
+        /*sqrt_alibi=*/false,
+        /*alibi_slopes=*/std::nullopt,
+        /*sinks=*/std::nullopt,
+        lse);
+}
+
+// Decode attention via xllm_paged_attention
+torch::Tensor ix_decode_attention(
+    torch::Tensor output,          // (num_seqs, num_heads, head_dim)
+    torch::Tensor query,           // (num_seqs, num_heads, head_dim)
+    torch::Tensor key_cache,       // (num_blocks, num_kv_heads, block_size, head_dim)
+    torch::Tensor value_cache,     // (num_blocks, num_kv_heads, block_size, head_dim)
+    int64_t num_kv_heads,
+    double scale,
+    torch::Tensor block_tables,    // (num_seqs, max_blocks)
+    torch::Tensor seq_lens,        // (num_seqs,)
+    int64_t block_size,
+    int64_t max_context_len) {
+
+    return ixformer::infer::xllm_paged_attention(
+        output, query, key_cache, value_cache,
+        num_kv_heads, scale,
+        block_tables, seq_lens,
+        block_size, max_context_len,
+        /*alibi_slopes=*/std::nullopt,
+        /*causal=*/true,
+        /*window_left=*/-1,
+        /*window_right=*/-1,
+        /*softcap=*/0.0,
+        /*enable_cuda_graph=*/false,
+        /*use_sqrt_alibi=*/false,
+        /*sinks=*/std::nullopt);
+}
+
+// Fused linear (matmul + optional activation)
+// act_type: 0=none, 1=silu, 2=gelu, 3=gelu_tanh
+torch::Tensor ix_linear(
+    torch::Tensor input,
+    torch::Tensor weight,
+    int64_t act_type) {
+    return ixformer::infer::ixformer_linear(
+        input, weight, act_type,
+        /*bias=*/std::nullopt,
+        /*out=*/std::nullopt,
+        /*persistent=*/std::nullopt);
+}
+
+// Fused residual + RMS norm
+// Port from: xllm/xllm/core/kernels/ilu/norm.cpp residual_layer_norm()
+std::tuple<torch::Tensor, torch::Tensor> ix_residual_rms_norm(
+    torch::Tensor input,
+    torch::Tensor residual,
+    torch::Tensor weight,
+    double eps) {
+    auto output = torch::zeros_like(input);
+    auto residual_output = torch::zeros_like(input);
+
+    ixformer::infer::residual_rms_norm(
+        input, residual, weight, output, residual_output,
+        /*fused_bias=*/std::nullopt,
+        /*alpha=*/1.0,
+        eps,
+        /*is_post=*/false);
+
+    return std::make_tuple(output, residual_output);
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("prefill_attention", &ix_prefill_attention,
+          "Fused prefill attention via ixformer flash_attn_unpad_with_block_tables",
+          py::arg("query"), py::arg("key_cache"), py::arg("value_cache"),
+          py::arg("output"), py::arg("block_tables"),
+          py::arg("cu_seq_q"), py::arg("cu_seq_k"),
+          py::arg("max_query_len"), py::arg("max_seq_len"),
+          py::arg("scale"),
+          py::arg("is_causal") = true,
+          py::arg("window_left") = -1,
+          py::arg("window_right") = -1);
+
+    m.def("decode_attention", &ix_decode_attention,
+          "Paged decode attention via ixformer xllm_paged_attention",
+          py::arg("output"), py::arg("query"),
+          py::arg("key_cache"), py::arg("value_cache"),
+          py::arg("num_kv_heads"), py::arg("scale"),
+          py::arg("block_tables"), py::arg("seq_lens"),
+          py::arg("block_size"), py::arg("max_context_len"));
+
+    m.def("linear", &ix_linear,
+          "Fused linear via ixformer (matmul + optional activation)",
+          py::arg("input"), py::arg("weight"), py::arg("act_type") = 0);
+
+    m.def("residual_rms_norm", &ix_residual_rms_norm,
+          "Fused residual + RMS norm via ixformer",
+          py::arg("input"), py::arg("residual"),
+          py::arg("weight"), py::arg("eps") = 1e-6);
+}

ex_engine/csrc/ix_full_bridge_v2.cpp

@@ -1,21 +1,24 @@
-// ix_full_bridge_v2.cpp — Complete bridge to ALL ixformer::infer C++ functions
+// ix_full_bridge_v2.cpp — Bridge to ixformer C++ functions + MoE pipeline
 //
-// Base image has ixformer::infer namespace with 14 functions.
-// Previous ix_full_bridge.cpp only bridged 4 (silu_and_mul, rms_norm,
-// fused_add_rms_norm, linear). This file bridges ALL 14.
+// Forward declarations use REAL symbols from nm -D symbol dumps:
+//   _ixformer_torch.so → namespace ixformer_torch_ext (7 functions)
+//   moe_ops_impl.cu    → namespace ixformer::infer    (5 MoE functions, self-compiled)
 //
-// The base image's _ixformer_torch.cpython-310.so and libixformer.so
-// export these symbols in the ixformer::infer namespace (confirmed by nm -D).
+// Symbol dump verified:
+//   ixformer_torch_ext::silu_and_mul_forward(at::Tensor&, at::Tensor&)
+//   ixformer_torch_ext::rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double)
+//   ixformer_torch_ext::fused_add_rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double, double)
+//   ixformer_torch_ext::ixformer_linear(at::Tensor&, at::Tensor&, c10::optional<at::Tensor>, c10::optional<at::Tensor>)
+//   ixformer_torch_ext::ixformer_linear_ex(at::Tensor&, at::Tensor&, c10::optional<at::Tensor>)
+//   ixformer_torch_ext::vllm_rotary_embedding_neox(at::Tensor&, at::Tensor&, at::Tensor&, long, at::Tensor&, long, bool)
+//   ixformer_torch_ext::vllm_cache_ops_reshape_and_cache(at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, long, long)
+//   ixformer_torch_ext::vllm_single_query_cached_kv_attention(13 params — see below)
 //
-// Compile:
-//   torch.utils.cpp_extension.load(
-//     name="ix_full_bridge_v2",
-//     sources=["ix_full_bridge_v2.cpp"],
-//     extra_ldflags=[<all ixformer .so files>, "-Wl,-rpath,..."],
-//     extra_cflags=["-O2", "-std=c++17"],
-//   )
-//
-// Upstream reference: xllm_latest/core/kernels/ilu/ixformer.h
+// NOT available in any .so (confirmed by nm -D on all 4 .so files):
+//   ixinfer_flash_attn_unpad_with_block_tables — DOES NOT EXIST
+//   xllm_paged_attention — DOES NOT EXIST
+//   topk_softmax, moe_w16a16_group_gemm, etc — NOT in libixformer.so
+//     (provided by moe_ops_impl.cu instead)
 
 #include <torch/extension.h>
 #include <optional>
@@ -24,103 +27,68 @@
 #include <vector>
 
 // ============================================================================
-// Forward declarations — ixformer::infer namespace from base image .so
-// Signatures EXACTLY match upstream_ref/xllm_latest/core/kernels/ilu/ixformer.h
+// Forward declarations — ixformer_torch_ext namespace from _ixformer_torch.so
+// Signatures EXACTLY match nm -D | c++filt output
+// ============================================================================
+namespace ixformer_torch_ext {
+
+// silu_and_mul_forward(at::Tensor&, at::Tensor&)
+void silu_and_mul_forward(at::Tensor& input, at::Tensor& output);
+
+// rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double)
+// Real ixformer signature order: (input, weight, output, eps)
+void rms_norm_forward(at::Tensor& input, at::Tensor& weight,
+                      at::Tensor& output, double eps);
+
+// fused_add_rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double, double)
+void fused_add_rms_norm_forward(at::Tensor& input, at::Tensor& residual,
+                                at::Tensor& weight, double eps, double alpha);
+
+// ixformer_linear(at::Tensor&, at::Tensor&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&)
+at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight,
+                           c10::optional<at::Tensor> const& bias,
+                           c10::optional<at::Tensor> const& out);
+
+// ixformer_linear_ex(at::Tensor&, at::Tensor&, c10::optional<at::Tensor> const&)
+at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight,
+                              c10::optional<at::Tensor> const& bias);
+
+// vllm_rotary_embedding_neox(at::Tensor&, at::Tensor&, at::Tensor&, long, at::Tensor&, long, bool)
+void vllm_rotary_embedding_neox(at::Tensor& positions, at::Tensor& query,
+                                at::Tensor& key, int64_t head_size,
+                                at::Tensor& cos_sin_cache,
+                                int64_t max_position, bool is_neox);
+
+// vllm_cache_ops_reshape_and_cache(at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, long, long)
+void vllm_cache_ops_reshape_and_cache(at::Tensor& key, at::Tensor& value,
+                                      at::Tensor& key_cache,
+                                      at::Tensor& value_cache,
+                                      at::Tensor& slot_mapping,
+                                      int64_t key_token_stride,
+                                      int64_t value_token_stride);
+
+// vllm_single_query_cached_kv_attention(at::Tensor& x13)
+// Full signature from nm -D:
+//   (at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&,
+//    double, at::Tensor&, at::Tensor&, long, long, long, bool,
+//    c10::optional<at::Tensor> const&)
+void vllm_single_query_cached_kv_attention(
+    at::Tensor& output, at::Tensor& query,
+    at::Tensor& key_cache, at::Tensor& value_cache,
+    at::Tensor& head_mapping, double scale,
+    at::Tensor& block_tables, at::Tensor& context_lens,
+    int64_t block_size, int64_t max_context_len, int64_t num_kv_heads,
+    bool is_neox,
+    c10::optional<at::Tensor> const& alibi_slopes);
+
+}  // namespace ixformer_torch_ext
+
+// ============================================================================
+// Forward declarations — ixformer::infer namespace from moe_ops_impl.cu
+// These 5 MoE functions are compiled from our own CUDA code, NOT from .so
 // ============================================================================
 namespace ixformer { namespace infer {
 
-// --- Attention ---
-torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
-    torch::Tensor& query,
-    torch::Tensor& key_cache,
-    torch::Tensor& value_cache,
-    torch::Tensor& out,
-    torch::Tensor& block_tables,
-    torch::Tensor& cu_seq_q,
-    torch::Tensor& cu_seq_k,
-    int64_t max_seq_q,
-    int64_t max_seq_k,
-    bool is_causal,
-    int64_t window_left,
-    int64_t window_right,
-    double scale,
-    double softcap,
-    bool sqrt_alibi,
-    const std::optional<torch::Tensor>& alibi_slopes,
-    const std::optional<torch::Tensor>& sinks,
-    std::optional<torch::Tensor>& lse);
-
-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);
-
-// --- Activation ---
-void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
-
-// --- Linear ---
-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);
-
-torch::Tensor ixformer_linear_ex(torch::Tensor& input,
-                                 torch::Tensor& weight,
-                                 const c10::optional<torch::Tensor>& bias,
-                                 const c10::optional<torch::Tensor>& out);
-
-// --- Cache ---
-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);
-
-// --- RoPE ---
-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);
-
-// --- Norm ---
-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);
-
-void rms_norm(torch::Tensor& input,
-              torch::Tensor& weight,
-              torch::Tensor& output,
-              const std::optional<torch::Tensor>& fused_bias,
-              double eps);
-
-// --- MoE ---
 void topk_softmax(torch::Tensor& topk_weights,
                   torch::Tensor& topk_indices,
                   torch::Tensor& token_expert_indices,
@@ -132,9 +100,9 @@ void moe_compute_token_index_api(
     torch::Tensor& src_dst,
     torch::Tensor& dst_src,
     torch::Tensor& expert_sizes_gpu,
-    const c10::optional<torch::Tensor>& expert_mask,
-    const c10::optional<torch::Tensor>& expert_sizes_cpu,
-    const c10::optional<torch::Tensor>& expand_tokens_gpu,
+    const std::optional<torch::Tensor>& expert_mask,
+    const std::optional<torch::Tensor>& expert_sizes_cpu,
+    const std::optional<torch::Tensor>& expand_tokens_gpu,
     int64_t start_expert_id,
     int64_t end_expert_id,
     int64_t num_experts);
@@ -142,7 +110,7 @@ void moe_compute_token_index_api(
 void moe_expand_input(torch::Tensor outputs,
                       torch::Tensor inputs,
                       torch::Tensor dst_to_src,
-                      const c10::optional<torch::Tensor>& src_to_dst,
+                      const std::optional<torch::Tensor>& src_to_dst,
                       int64_t dst_tokens,
                       int64_t expand_factor);
 
@@ -150,52 +118,47 @@ void moe_w16a16_group_gemm(torch::Tensor output,
                            torch::Tensor inputs,
                            torch::Tensor weights,
                            torch::Tensor tokens_per_experts,
-                           const c10::optional<torch::Tensor>& dst_to_src,
-                           const c10::optional<torch::Tensor>& bias,
+                           const std::optional<torch::Tensor>& dst_to_src,
+                           const std::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 c10::optional<torch::Tensor>& mul_weight,
-                           const c10::optional<torch::Tensor>& mask,
-                           const c10::optional<torch::Tensor>& extra_residual,
+                           const std::optional<torch::Tensor>& mul_weight,
+                           const std::optional<torch::Tensor>& mask,
+                           const std::optional<torch::Tensor>& extra_residual,
                            double scaling_factor);
 
 }}  // namespace ixformer::infer
 
 
 // ============================================================================
-// Python wrappers — thin wrappers that match ix_bridge.py's expected API
+// Python wrappers — thin wrappers matching ix_bridge.py's expected API
 // ============================================================================
 
 // --- silu_and_mul ---
 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);
+    ixformer_torch_ext::silu_and_mul_forward(input, output);
     return output;
 }
 
 // --- rms_norm ---
 void ix_rms_norm(torch::Tensor output, torch::Tensor input,
                  torch::Tensor weight, double eps) {
-    ixformer::infer::rms_norm(input, weight, output,
-                              /*fused_bias=*/std::nullopt, eps);
+    // pybind receives (output, input, weight, eps)
+    // ixformer expects (input, weight, output, eps)
+    ixformer_torch_ext::rms_norm_forward(input, weight, output, eps);
 }
 
 // --- fused_add_rms_norm ---
-// residual_rms_norm does: output = rms_norm(input + alpha*residual, weight, eps)
-//                         residual_output = input + alpha*residual
 void ix_fused_add_rms_norm(torch::Tensor input, torch::Tensor residual,
-                           torch::Tensor weight, torch::Tensor output,
-                           torch::Tensor residual_output, double eps) {
-    ixformer::infer::residual_rms_norm(input, residual, weight,
-                                       output, residual_output,
-                                       /*fused_bias=*/std::nullopt,
-                                       /*alpha=*/1.0, eps,
-                                       /*is_post=*/false);
+                           torch::Tensor weight, double eps) {
+    ixformer_torch_ext::fused_add_rms_norm_forward(
+        input, residual, weight, eps, /*alpha=*/1.0);
 }
 
 // --- linear ---
@@ -204,74 +167,56 @@ torch::Tensor ix_linear(torch::Tensor input, torch::Tensor weight,
     auto input_2d = input.view({-1, input.size(-1)});
     int64_t m = input_2d.size(0);
     if (m <= 1 && !bias.has_value()) {
-        return ixformer::infer::ixformer_linear_ex(
-            input, weight, bias, /*out=*/c10::optional<torch::Tensor>());
+        return ixformer_torch_ext::ixformer_linear_ex(input, weight, bias);
     }
-    return ixformer::infer::ixformer_linear(
-        input, weight, /*act_type=*/0, bias,
-        /*out=*/std::nullopt, /*persistent=*/std::nullopt);
+    return ixformer_torch_ext::ixformer_linear(
+        input, weight, bias, /*out=*/c10::optional<at::Tensor>());
 }
 
 // --- rotary_embedding ---
 void ix_rotary_embedding(torch::Tensor positions, torch::Tensor query,
                          torch::Tensor key, int64_t head_size,
                          torch::Tensor cos_sin_cache, bool is_neox) {
-    ixformer::infer::xllm_rotary_embedding(
-        positions, query, key, head_size, cos_sin_cache, is_neox);
+    int64_t max_position = cos_sin_cache.size(0);
+    ixformer_torch_ext::vllm_rotary_embedding_neox(
+        positions, query, key, head_size, cos_sin_cache, max_position, is_neox);
 }
 
 // --- reshape_and_cache ---
 void ix_reshape_and_cache(torch::Tensor key, torch::Tensor value,
                           torch::Tensor key_cache, torch::Tensor value_cache,
                           torch::Tensor slot_mapping) {
-    // token stride = product of dims after dim 0 for key/value
-    // key shape: [num_tokens, num_heads, head_dim]
     int64_t key_token_stride = 1;
     for (int i = 1; i < key.dim(); i++) key_token_stride *= key.size(i);
     int64_t value_token_stride = 1;
     for (int i = 1; i < value.dim(); i++) value_token_stride *= value.size(i);
 
-    ixformer::infer::xllm_reshape_and_cache(
+    ixformer_torch_ext::vllm_cache_ops_reshape_and_cache(
         key, value, key_cache, value_cache, slot_mapping,
         key_token_stride, value_token_stride);
 }
 
-// --- paged_attention (decode) ---
-torch::Tensor ix_paged_attention(
+// --- paged_attention (decode only — no prefill available in .so) ---
+void ix_paged_attention(
     torch::Tensor output, torch::Tensor query,
     torch::Tensor key_cache, torch::Tensor value_cache,
-    int64_t num_kv_heads, double scale,
+    torch::Tensor head_mapping, double scale,
     torch::Tensor block_tables, torch::Tensor context_lens,
-    int64_t block_size, int64_t max_context_len,
+    int64_t block_size, int64_t max_context_len, int64_t num_kv_heads,
     const c10::optional<torch::Tensor>& alibi_slopes) {
-    return ixformer::infer::xllm_paged_attention(
+    ixformer_torch_ext::vllm_single_query_cached_kv_attention(
         output, query, key_cache, value_cache,
-        num_kv_heads, scale, block_tables, context_lens,
-        block_size, max_context_len, alibi_slopes,
-        /*causal=*/true, /*window_left=*/-1, /*window_right=*/-1,
-        /*softcap=*/0.0, /*enable_cuda_graph=*/false,
-        /*use_sqrt_alibi=*/false, /*sinks=*/std::nullopt);
+        head_mapping, scale, block_tables, context_lens,
+        block_size, max_context_len, num_kv_heads,
+        /*is_neox=*/true, alibi_slopes);
 }
 
-// --- flash_attn_prefill ---
-torch::Tensor ix_flash_attn_prefill(
-    torch::Tensor query, torch::Tensor key_cache, torch::Tensor value_cache,
-    torch::Tensor output, torch::Tensor block_tables,
-    torch::Tensor cu_seq_q, torch::Tensor cu_seq_k,
-    int64_t max_query_len, int64_t max_seq_len,
-    double scale, bool is_causal,
-    int64_t window_left, int64_t window_right) {
-    std::optional<torch::Tensor> lse = std::nullopt;
-    return ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables(
-        query, key_cache, value_cache, output, block_tables,
-        cu_seq_q, cu_seq_k, max_query_len, max_seq_len,
-        is_causal, window_left, window_right, scale,
-        /*softcap=*/0.0, /*sqrt_alibi=*/false,
-        /*alibi_slopes=*/std::nullopt, /*sinks=*/std::nullopt, lse);
-}
 
-// --- MoE: topk_softmax ---
-// Returns (topk_weights, topk_ids, token_expert_indices)
+// ============================================================================
+// MoE wrappers — call moe_ops_impl.cu implementations
+// ============================================================================
+
+// --- topk_softmax ---
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
 ix_topk_softmax(torch::Tensor gating_output, int64_t topk, bool renormalize) {
     int64_t num_tokens = gating_output.size(0);
@@ -289,8 +234,7 @@ ix_topk_softmax(torch::Tensor gating_output, int64_t topk, bool renormalize) {
     return std::make_tuple(topk_weights, topk_ids, token_expert_indices);
 }
 
-// --- MoE: moe_gen_idx ---
-// Equivalent to xllm::kernel::ilu::moe_gen_idx
+// --- 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()});
@@ -299,9 +243,9 @@ ix_moe_gen_idx(torch::Tensor expert_id, int64_t expert_num) {
 
     ixformer::infer::moe_compute_token_index_api(
         expert_id, src_dst, dst_src, expert_sizes_gpu,
-        /*expert_mask=*/c10::nullopt,
-        /*expert_sizes_cpu=*/c10::nullopt,
-        /*expand_tokens_gpu=*/c10::nullopt,
+        /*expert_mask=*/std::nullopt,
+        /*expert_sizes_cpu=*/std::nullopt,
+        /*expand_tokens_gpu=*/std::nullopt,
         /*start_expert_id=*/0,
         /*end_expert_id=*/expert_num,
         /*num_experts=*/expert_num);
@@ -310,7 +254,7 @@ ix_moe_gen_idx(torch::Tensor expert_id, int64_t expert_num) {
     return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_cumsum};
 }
 
-// --- MoE: moe_expand_input ---
+// --- moe_expand_input ---
 torch::Tensor ix_moe_expand_input(torch::Tensor input,
                                   torch::Tensor gather_index,
                                   torch::Tensor combine_idx,
@@ -322,43 +266,41 @@ torch::Tensor ix_moe_expand_input(torch::Tensor input,
     return output;
 }
 
-// --- MoE: group_gemm ---
+// --- group_gemm ---
 torch::Tensor ix_group_gemm(torch::Tensor inputs, torch::Tensor weights,
                             torch::Tensor tokens_per_experts,
                             int64_t output_n) {
     int64_t total_tokens = inputs.size(0);
     auto output = inputs.new_empty({total_tokens, output_n});
+    int64_t gemm_output_n = tokens_per_experts.sum().item<int64_t>();
     ixformer::infer::moe_w16a16_group_gemm(
         output, inputs, weights, tokens_per_experts,
-        /*dst_to_src=*/c10::nullopt,
-        /*bias=*/c10::nullopt,
-        /*format=*/"default",
+        /*dst_to_src=*/std::nullopt,
+        /*bias=*/std::nullopt,
+        /*format=*/"TN",
         /*persistent=*/0,
-        output_n);
+        gemm_output_n);
     return output;
 }
 
-// --- MoE: moe_combine_result ---
+// --- moe_combine_result ---
 torch::Tensor ix_moe_combine_result(torch::Tensor input, torch::Tensor weight) {
-    // input: [T*topk, H], weight: [T, topk]
     auto input_3d = input.view({-1, weight.size(1), input.size(1)});
     auto output = input.new_empty({input_3d.size(0), input_3d.size(2)});
     ixformer::infer::moe_output_reduce_sum(
         output, input_3d, weight,
-        /*mask=*/c10::nullopt,
-        /*extra_residual=*/c10::nullopt,
+        /*mask=*/std::nullopt,
+        /*extra_residual=*/std::nullopt,
         /*scaling_factor=*/1.0);
     return output;
 }
 
-// --- MoE: fused_moe_forward (7-step pipeline) ---
-// This is the full fused MoE forward: topk → gen_idx → expand → gemm(w13) →
-// silu_mul → gemm(w2) → combine
+// --- fused_moe_forward (7-step pipeline) ---
 torch::Tensor ix_fused_moe_forward(
     torch::Tensor hidden_states,
     torch::Tensor router_logits,
-    torch::Tensor w13,       // [num_experts, 2*intermediate, hidden]
-    torch::Tensor w2,        // [num_experts, hidden, intermediate]
+    torch::Tensor w13,
+    torch::Tensor w2,
     int64_t topk,
     int64_t num_experts,
     bool renormalize) {
@@ -383,7 +325,7 @@ torch::Tensor ix_fused_moe_forward(
 
     // Step 4: group_gemm (w13: gate_up projection)
     int64_t intermediate_2x = w13.size(1);
-    auto gate_up = ix_group_gemm(expanded, w13.view({-1, w13.size(2)}),
+    auto gate_up = ix_group_gemm(expanded, w13,
                                  expert_sizes_gpu, intermediate_2x);
 
     // Step 5: silu_and_mul
@@ -391,7 +333,7 @@ torch::Tensor ix_fused_moe_forward(
 
     // Step 6: group_gemm (w2: down projection)
     int64_t hidden_size = w2.size(1);
-    auto down = ix_group_gemm(activated, w2.view({-1, w2.size(2)}),
+    auto down = ix_group_gemm(activated, w2,
                               expert_sizes_gpu, hidden_size);
 
     // Step 7: moe_combine_result
@@ -402,50 +344,48 @@ torch::Tensor ix_fused_moe_forward(
 
 
 // ============================================================================
-// Module registration — ALL 14 functions + fused pipeline
+// Module registration
 // ============================================================================
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     // Activation
     m.def("silu_and_mul", &ix_silu_and_mul,
-          "Fused SiLU+mul activation via ixformer::infer");
+          "Fused SiLU+mul via ixformer_torch_ext");
 
     // Norm
     m.def("rms_norm", &ix_rms_norm,
-          "RMSNorm via ixformer::infer");
+          "RMSNorm via ixformer_torch_ext");
     m.def("fused_add_rms_norm", &ix_fused_add_rms_norm,
-          "Residual + RMSNorm via ixformer::infer");
+          "Residual + RMSNorm via ixformer_torch_ext");
 
     // Linear
     m.def("linear", &ix_linear,
-          "GEMM via ixformer::infer (linear/linear_ex)");
+          "GEMM via ixformer_torch_ext");
 
     // RoPE
     m.def("rotary_embedding", &ix_rotary_embedding,
-          "Rotary position embedding via ixformer::infer");
+          "Rotary embedding via ixformer_torch_ext");
 
     // Cache
     m.def("reshape_and_cache", &ix_reshape_and_cache,
-          "KV cache reshape+store via ixformer::infer");
+          "KV cache reshape+store via ixformer_torch_ext");
 
-    // Attention
+    // Attention (decode only)
     m.def("paged_attention", &ix_paged_attention,
-          "Paged attention decode via ixformer::infer");
-    m.def("flash_attn_prefill", &ix_flash_attn_prefill,
-          "Flash attention prefill via ixformer::infer");
+          "Paged attention decode via ixformer_torch_ext");
 
-    // MoE (individual steps)
+    // MoE (individual steps — from moe_ops_impl.cu)
     m.def("topk_softmax", &ix_topk_softmax,
-          "MoE topk+softmax routing via ixformer::infer");
+          "MoE topk+softmax routing");
     m.def("moe_gen_idx", &ix_moe_gen_idx,
-          "MoE compute token index via ixformer::infer");
+          "MoE compute token index");
     m.def("moe_expand_input", &ix_moe_expand_input,
-          "MoE expand input for expert dispatch via ixformer::infer");
+          "MoE expand input for expert dispatch");
     m.def("group_gemm", &ix_group_gemm,
-          "MoE grouped GEMM via ixformer::infer");
+          "MoE grouped GEMM via cuinferCustomGemm");
     m.def("moe_combine_result", &ix_moe_combine_result,
-          "MoE output reduce sum via ixformer::infer");
+          "MoE output reduce sum");
 
     // MoE (fused 7-step pipeline)
     m.def("fused_moe_forward", &ix_fused_moe_forward,
-          "Complete fused MoE forward (7-step pipeline) via ixformer::infer");
-}
+          "Complete fused MoE forward (7-step pipeline)");
+}

ex_engine/csrc/ix_moe_bridge.cpp

@@ -149,7 +149,7 @@ torch::Tensor ix_group_gemm(
       output, inputs, weights, token_count,
       /*dst_to_src=*/kNoneTensor,
       /*bias=*/kNoneTensor,
-      /*format=*/"NT",
+      /*format=*/"TN",
       /*persistent=*/0,
       /*output_n=*/output_n);
   return output;

ex_engine/csrc/moe_expert_gemm.cpp

@@ -0,0 +1,180 @@
+// moe_expert_gemm.cpp — MoE expert GEMM dispatch
+//
+// Replaces the Python for-loop over experts with a C++ loop calling
+// ixformer_linear (via base image's _ixformer_torch.so).
+//
+// Why this works:
+//   1. Eliminates Python interpreter overhead per expert (~0.5ms × 64 experts)
+//   2. Eliminates PyTorch dispatcher overhead per F.linear call
+//   3. Uses the same ixformer GEMM kernel that the base image uses
+//   4. No new dependencies — links against the same .so as ix_full_bridge
+//
+// For decode (single token, top_k=8 experts):
+//   Python: 8 × F.linear → 8 × Python dispatch → 8 × CUDA kernel
+//   This:   1 × Python call → 8 × C++ ixformer_linear → 8 × CUDA kernel
+//   Savings: ~4ms → ~0.5ms (eliminate 7 Python round-trips)
+//
+// For prefill (many tokens, up to 64 experts):
+//   Python: for eid in 64: F.linear(tokens[eid], w[eid])
+//   This:   1 × Python call → C++ loop: 64 × ixformer_linear
+//   Savings: ~32ms → ~4ms
+//
+// Future: replace C++ loop with cublasGemmBatchedEx for true batched GEMM
+
+#include <torch/extension.h>
+#include <optional>
+#include <vector>
+
+// ============================================================================
+// Forward declarations — from base image _ixformer_torch.cpython-310.so
+// ============================================================================
+namespace ixformer_torch_ext {
+
+at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight,
+                           const c10::optional<at::Tensor>& bias,
+                           const c10::optional<at::Tensor>& out);
+
+at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight,
+                              const c10::optional<at::Tensor>& bias);
+
+void silu_and_mul_forward(at::Tensor& input, at::Tensor& output);
+
+}  // namespace ixformer_torch_ext
+
+
+// ============================================================================
+// Decode path: single token, top_k experts
+// ============================================================================
+// Input:  hidden (1, H), w13 (E, 2*I, H), w2 (E, H, I), expert_ids (K,), weights (K,)
+// Output: (1, H)
+//
+// Steps per expert:
+//   1. gate_up = ixformer_linear(hidden, w13[eid])   → (1, 2*I)
+//   2. act = silu_and_mul(gate_up)                    → (1, I)
+//   3. expert_out = ixformer_linear(act, w2[eid])     → (1, H)
+//   4. accumulate: out += weight[k] * expert_out
+
+torch::Tensor moe_decode_experts(
+    torch::Tensor hidden,       // (1, H)
+    torch::Tensor w13,          // (num_experts, 2*inter, H)
+    torch::Tensor w2,           // (num_experts, H, inter)
+    torch::Tensor expert_ids,   // (top_k,)  int64
+    torch::Tensor expert_weights  // (top_k,) fp16/fp32
+) {
+    int64_t top_k = expert_ids.size(0);
+    int64_t H = hidden.size(-1);
+    int64_t inter2 = w13.size(1);  // 2 * intermediate
+    int64_t inter = inter2 / 2;
+
+    auto out = torch::zeros({1, H}, hidden.options());
+    c10::optional<at::Tensor> no_bias;
+
+    for (int64_t k = 0; k < top_k; ++k) {
+        int64_t eid = expert_ids[k].item<int64_t>();
+        float w = expert_weights[k].item<float>();
+
+        // w13[eid] shape: (2*I, H) — use as weight for linear
+        auto w13_e = w13[eid];  // (2*I, H)
+        auto w2_e = w2[eid];    // (H, I)
+
+        // gate_up = hidden @ w13_e^T → (1, 2*I)
+        auto gate_up = ixformer_torch_ext::ixformer_linear(
+            hidden, w13_e, no_bias, c10::optional<at::Tensor>());
+
+        // silu_and_mul: (1, 2*I) → (1, I)
+        auto act = torch::empty({1, inter}, hidden.options());
+        ixformer_torch_ext::silu_and_mul_forward(gate_up, act);
+
+        // expert_out = act @ w2_e^T → (1, H)
+        auto expert_out = ixformer_torch_ext::ixformer_linear(
+            act, w2_e, no_bias, c10::optional<at::Tensor>());
+
+        // accumulate
+        out.add_(expert_out, w);
+    }
+
+    return out;
+}
+
+
+// ============================================================================
+// Prefill path: multiple tokens, grouped by expert
+// ============================================================================
+// Input:  hidden (T, H), w13 (E, 2*I, H), w2 (E, H, I),
+//         sorted_token_ids (T*K,), sorted_weights (T*K,), expert_counts list
+// Output: (T, H)
+//
+// For each expert with count > 0:
+//   tokens = hidden[sorted_token_ids[start:end]]
+//   gate_up = ixformer_linear(tokens, w13[eid])
+//   act = silu_and_mul(gate_up)
+//   expert_out = ixformer_linear(act, w2[eid])
+//   out[token_ids] += expert_out * weights
+
+torch::Tensor moe_prefill_experts(
+    torch::Tensor hidden,           // (T, H)
+    torch::Tensor w13,              // (E, 2*I, H)
+    torch::Tensor w2,               // (E, H, I)
+    torch::Tensor sorted_token_ids, // (T*K,) int64
+    torch::Tensor sorted_weights,   // (T*K,) fp16/fp32
+    torch::Tensor expert_counts     // (E,) int64
+) {
+    int64_t T = hidden.size(0);
+    int64_t H = hidden.size(-1);
+    int64_t inter2 = w13.size(1);
+    int64_t inter = inter2 / 2;
+    int64_t E = expert_counts.size(0);
+
+    auto out = torch::zeros({T, H}, hidden.options());
+    c10::optional<at::Tensor> no_bias;
+
+    int64_t start = 0;
+    for (int64_t eid = 0; eid < E; ++eid) {
+        int64_t count = expert_counts[eid].item<int64_t>();
+        if (count == 0) continue;
+        int64_t end = start + count;
+
+        auto tok_ids = sorted_token_ids.slice(0, start, end);   // (count,)
+        auto tokens = hidden.index_select(0, tok_ids);           // (count, H)
+        auto weights = sorted_weights.slice(0, start, end);      // (count,)
+
+        auto w13_e = w13[eid];  // (2*I, H)
+        auto w2_e = w2[eid];    // (H, I)
+
+        // FC1: gate_up = tokens @ w13_e^T → (count, 2*I)
+        auto gate_up = ixformer_torch_ext::ixformer_linear(
+            tokens, w13_e, no_bias, c10::optional<at::Tensor>());
+
+        // SiLU and mul: (count, 2*I) → (count, I)
+        auto act = torch::empty({count, inter}, hidden.options());
+        ixformer_torch_ext::silu_and_mul_forward(gate_up, act);
+
+        // FC2: expert_out = act @ w2_e^T → (count, H)
+        auto expert_out = ixformer_torch_ext::ixformer_linear(
+            act, w2_e, no_bias, c10::optional<at::Tensor>());
+
+        // Weighted accumulate: out[tok_ids] += expert_out * weights
+        auto weighted = expert_out * weights.unsqueeze(-1);
+        out.index_add_(0, tok_ids, weighted.to(out.dtype()));
+
+        start = end;
+    }
+
+    return out;
+}
+
+
+// ============================================================================
+// Module registration
+// ============================================================================
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("moe_decode_experts", &moe_decode_experts,
+          "MoE decode: C++ loop over top_k experts using ixformer_linear",
+          py::arg("hidden"), py::arg("w13"), py::arg("w2"),
+          py::arg("expert_ids"), py::arg("expert_weights"));
+    m.def("moe_prefill_experts", &moe_prefill_experts,
+          "MoE prefill: C++ loop over experts using ixformer_linear",
+          py::arg("hidden"), py::arg("w13"), py::arg("w2"),
+          py::arg("sorted_token_ids"), py::arg("sorted_weights"),
+          py::arg("expert_counts"));
+}

ex_engine/csrc/moe_ops_impl.cu

@@ -0,0 +1,502 @@
+// moe_ops_impl.cu — Implement the 5 missing MoE functions
+//
+// These functions are declared in ixformer.h (from xllm upstream)
+// but NOT present in the base image's libixformer.so.
+//
+// We implement them using available primitives:
+//   - cuinferCustomGemm (from libcuinfer.so) for group_gemm
+//   - Pure CUDA kernels for topk_softmax, moe_compute_index, expand, combine
+//   - ixformer::functions::cuinfer_gemm (from libixformer.so) as fallback
+//
+// Reference AST chain:
+//   xllm/core/kernels/ilu/fused_moe.cpp  → calls these 5 functions
+//   xllm/core/kernels/ilu/group_gemm.cpp → calls moe_w16a16_group_gemm
+//   xllm/core/kernels/ilu/ixformer.h     → declares them in ixformer::infer
+//
+// We provide them in the SAME namespace so ix_full_bridge_v2.cpp links cleanly.
+
+#include <torch/extension.h>
+#include <c10/cuda/CUDAStream.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <optional>
+#include <vector>
+#include <numeric>
+
+// ============================================================================
+// Forward-declare cuinfer C API (from libcuinfer.so, confirmed in symbol dump)
+// ============================================================================
+extern "C" {
+
+typedef struct cuinferContext* cuinferHandle_t;
+typedef enum { CUINFER_STATUS_SUCCESS = 0 } cuinferStatus_t;
+typedef enum {
+    CUINFER_OP_TENSOR_OP_N = 0,
+    CUINFER_OP_TENSOR_OP_T = 1,
+} cuinferOperation_t;
+typedef enum {
+    CUINFER_GEMM_DEFAULT = 0,
+} cuinferGEMMCustomOption_t;
+typedef enum {
+    CUINFER_POINTER_MODE_HOST = 0,
+} cuinferPointerMode_t;
+
+cuinferStatus_t cuinferCreate(cuinferHandle_t* handle);
+cuinferStatus_t cuinferDestroy(cuinferHandle_t handle);
+cuinferStatus_t cuinferSetStream(cuinferHandle_t handle, cudaStream_t stream);
+
+cuinferStatus_t cuinferCustomGemm(
+    cuinferHandle_t handle, cudaStream_t stream,
+    cuinferPointerMode_t ptrMode,
+    cuinferOperation_t transa, cuinferOperation_t transb,
+    int m, int n, int k,
+    const void* alpha,
+    const void* A, cudaDataType_t Atype, int lda, long long int strideA,
+    const void* B, cudaDataType_t Btype, int ldb, long long int strideB,
+    const void* beta,
+    void* C, cudaDataType_t Ctype, int ldc, long long int strideC,
+    int batchCount,
+    cudaDataType_t computeType, cudaDataType_t scaleType,
+    const void* customHostPtr, const void* customDevicePtr,
+    cuinferGEMMCustomOption_t customOption);
+
+}  // extern "C"
+
+
+// ============================================================================
+// Kernel 1: topk_softmax
+// Adapted from moe_topk_softmax_v3.cu (already working, 64-expert specialized)
+// ============================================================================
+
+// Qwen3.5-27B: 128 routed experts
+// Block size = 128 threads (1 thread per expert for ≤128 experts)
+static constexpr int MOE_MAX_EXPERTS = 128;
+static constexpr int MOE_BLOCK = 128;
+
+// All reductions use blockDim.x (dynamic block size, power-of-2)
+__device__ float smem_reduce_max(float val, float* smem) {
+    int tid = threadIdx.x;
+    smem[tid] = val;
+    __syncthreads();
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
+        __syncthreads();
+    }
+    return smem[0];
+}
+
+__device__ float smem_reduce_sum(float val, float* smem) {
+    int tid = threadIdx.x;
+    smem[tid] = val;
+    __syncthreads();
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) smem[tid] += smem[tid + s];
+        __syncthreads();
+    }
+    return smem[0];
+}
+
+__device__ void smem_argmax(float val, int idx, float* s_val, int* s_idx) {
+    int tid = threadIdx.x;
+    s_val[tid] = val;
+    s_idx[tid] = idx;
+    __syncthreads();
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s && s_val[tid + s] > s_val[tid]) {
+            s_val[tid] = s_val[tid + s];
+            s_idx[tid] = s_idx[tid + s];
+        }
+        __syncthreads();
+    }
+}
+
+__global__ void topk_softmax_kernel(
+    const float* __restrict__ input,
+    float* __restrict__ topk_weights,
+    int32_t* __restrict__ topk_indices,
+    int32_t* __restrict__ token_expert_indices,
+    int num_tokens, int num_experts, int topk, bool renormalize
+) {
+    int row = blockIdx.x;
+    if (row >= num_tokens) return;
+    int tid = threadIdx.x;
+
+    extern __shared__ char shared_buf[];
+    float* smem = (float*)shared_buf;
+    int* smem_idx = (int*)(smem + blockDim.x);
+
+    // num_experts passed via gridDim.y (encoded), or read from shared
+    // We use a separate parameter for clarity
+    float val = (tid < num_experts) ? input[row * num_experts + tid] : -1e30f;
+
+    // Softmax
+    float row_max = smem_reduce_max(val, smem);
+    val = (tid < num_experts) ? expf(val - row_max) : 0.0f;
+    float row_sum = smem_reduce_sum(val, smem);
+    val *= (1.0f / row_sum);
+
+    float* out_w = topk_weights + row * topk;
+    int32_t* out_idx = topk_indices + row * topk;
+    int32_t* out_src = token_expert_indices + row * topk;
+
+    float my_val = val;
+    float topk_sum = 0.0f;
+
+    for (int ki = 0; ki < topk; ki++) {
+        smem_argmax(my_val, tid, smem, smem_idx);
+        float winner_val = smem[0];
+        int winner_idx = smem_idx[0];
+        __syncthreads();
+
+        if (tid == 0) {
+            out_w[ki] = winner_val;
+            out_idx[ki] = winner_idx;
+            out_src[ki] = row;
+        }
+        topk_sum += winner_val;
+        if (tid == winner_idx) my_val = -1.0f;
+        __syncthreads();
+    }
+
+    if (renormalize && tid == 0) {
+        float inv = 1.0f / (topk_sum + 1e-8f);
+        for (int ki = 0; ki < topk; ki++)
+            out_w[ki] *= inv;
+    }
+}
+
+
+// ============================================================================
+// Kernel 2: moe_compute_token_index
+// Histogram + prefix sum + scatter — from xllm_kernels/cuda/moe_compute_index.cu
+// ============================================================================
+
+__global__ void histogram_kernel(
+    const int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ expert_sizes,
+    int num_elements, int num_experts
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < num_elements) {
+        int eid = expert_ids[idx];
+        if (eid >= 0 && eid < num_experts) {
+            atomicAdd(&expert_sizes[eid], 1);
+        }
+    }
+}
+
+__global__ void place_indices_kernel(
+    const int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ expert_offsets,  // will be atomicAdd'd
+    int32_t* __restrict__ src_dst,
+    int32_t* __restrict__ dst_src,
+    int num_elements
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < num_elements) {
+        int eid = expert_ids[idx];
+        int pos = atomicAdd(&expert_offsets[eid], 1);
+        src_dst[idx] = pos;  // where token idx goes in sorted order
+        dst_src[pos] = idx;  // reverse mapping
+    }
+}
+
+
+// ============================================================================
+// Kernel 3: moe_expand_input
+// Gather-based expand: output[i] = input[gather_index[i]]
+// ============================================================================
+
+template <typename scalar_t>
+__global__ void expand_input_kernel(
+    scalar_t* __restrict__ output,
+    const scalar_t* __restrict__ input,
+    const int32_t* __restrict__ dst_to_src,
+    int num_output_tokens, int hidden_size
+) {
+    int token = blockIdx.x;
+    if (token >= num_output_tokens) return;
+
+    int src_token = dst_to_src[token];
+    const scalar_t* src = input + (int64_t)src_token * hidden_size;
+    scalar_t* dst = output + (int64_t)token * hidden_size;
+
+    for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
+        dst[h] = src[h];
+    }
+}
+
+
+// ============================================================================
+// Kernel 4: moe_combine_result (weighted sum of expert outputs)
+// output[t] = sum_k( weight[t][k] * gemm2_output[flat_index(t,k)] )
+// ============================================================================
+
+template <typename scalar_t>
+__global__ void combine_result_kernel(
+    scalar_t* __restrict__ output,      // [N, H]
+    const scalar_t* __restrict__ input,  // [N*topk, H]
+    const float* __restrict__ weights,   // [N, topk]
+    int num_tokens, int topk, int hidden_size
+) {
+    int token = blockIdx.x;
+    if (token >= num_tokens) return;
+
+    for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
+        float acc = 0.0f;
+        for (int k = 0; k < topk; k++) {
+            int flat = token * topk + k;
+            float w = weights[token * topk + k];
+            acc += w * __half2float(input[flat * hidden_size + h]);
+        }
+        output[token * hidden_size + h] = __float2half(acc);
+    }
+}
+
+// Float specialization
+template <>
+__global__ void combine_result_kernel<float>(
+    float* __restrict__ output,
+    const float* __restrict__ input,
+    const float* __restrict__ weights,
+    int num_tokens, int topk, int hidden_size
+) {
+    int token = blockIdx.x;
+    if (token >= num_tokens) return;
+
+    for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
+        float acc = 0.0f;
+        for (int k = 0; k < topk; k++) {
+            int flat = token * topk + k;
+            float w = weights[token * topk + k];
+            acc += w * input[flat * hidden_size + h];
+        }
+        output[token * hidden_size + h] = acc;
+    }
+}
+
+
+// ============================================================================
+// C++ wrapper functions — ixformer::infer namespace
+// These provide the MISSING symbols that ix_full_bridge_v2.cpp needs.
+// ============================================================================
+
+namespace ixformer { namespace infer {
+
+void topk_softmax(
+    torch::Tensor& topk_weights,
+    torch::Tensor& topk_indices,
+    torch::Tensor& token_expert_indices,
+    torch::Tensor& gating_output,
+    bool renormalize
+) {
+    int num_tokens = gating_output.size(0);
+    int num_experts = gating_output.size(1);
+    int topk = topk_weights.size(1);
+    auto stream = c10::cuda::getCurrentCUDAStream();
+
+    auto input_f32 = gating_output.to(torch::kFloat32).contiguous();
+
+    // Block size must be >= num_experts, round up to next power of 2
+    int block_size = 1;
+    while (block_size < num_experts) block_size <<= 1;
+    TORCH_CHECK(block_size <= 1024, "Too many experts for topk kernel: ", num_experts);
+
+    size_t smem_bytes = block_size * (sizeof(float) + sizeof(int));
+    topk_softmax_kernel<<<num_tokens, block_size, smem_bytes, stream>>>(
+        input_f32.data_ptr<float>(),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int32_t>(),
+        token_expert_indices.data_ptr<int32_t>(),
+        num_tokens, num_experts, topk, renormalize);
+}
+
+void moe_compute_token_index_api(
+    torch::Tensor& topk_ids,
+    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,
+    int64_t start_expert_id,
+    int64_t end_expert_id,
+    int64_t num_experts
+) {
+    auto stream = c10::cuda::getCurrentCUDAStream();
+    int num_elements = topk_ids.numel();
+
+    // Zero expert_sizes
+    cudaMemsetAsync(expert_sizes_gpu.data_ptr<int32_t>(), 0,
+                    num_experts * sizeof(int32_t), stream);
+
+    // Phase 1: histogram
+    int blocks1 = (num_elements + 255) / 256;
+    histogram_kernel<<<blocks1, 256, 0, stream>>>(
+        topk_ids.data_ptr<int32_t>(),
+        expert_sizes_gpu.data_ptr<int32_t>(),
+        num_elements, num_experts);
+
+    // Phase 2: prefix sum for offsets (exclusive scan on GPU)
+    // Use a separate buffer for offsets, then reset for place_indices
+    auto expert_offsets = torch::zeros({num_experts}, topk_ids.options().dtype(torch::kInt32));
+    // Copy sizes → do exclusive scan on CPU (small: 64 experts)
+    auto sizes_cpu = expert_sizes_gpu.to(torch::kCPU);
+    auto offsets_cpu = torch::zeros({num_experts}, torch::dtype(torch::kInt32));
+    int32_t* s = sizes_cpu.data_ptr<int32_t>();
+    int32_t* o = offsets_cpu.data_ptr<int32_t>();
+    int32_t running = 0;
+    for (int i = 0; i < num_experts; i++) {
+        o[i] = running;
+        running += s[i];
+    }
+    expert_offsets = offsets_cpu.to(topk_ids.device());
+
+    // Phase 3: place indices
+    int blocks3 = (num_elements + 255) / 256;
+    place_indices_kernel<<<blocks3, 256, 0, stream>>>(
+        topk_ids.data_ptr<int32_t>(),
+        expert_offsets.data_ptr<int32_t>(),
+        src_dst.data_ptr<int32_t>(),
+        dst_src.data_ptr<int32_t>(),
+        num_elements);
+}
+
+void moe_expand_input(
+    torch::Tensor outputs,
+    torch::Tensor inputs,
+    torch::Tensor dst_to_src,
+    const std::optional<torch::Tensor>& src_to_dst,
+    int64_t dst_tokens,
+    int64_t expand_factor
+) {
+    auto stream = c10::cuda::getCurrentCUDAStream();
+    int hidden_size = inputs.size(1);
+    int block = std::min(hidden_size, 256);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(inputs.scalar_type(), "expand_input", [&] {
+        expand_input_kernel<scalar_t><<<dst_tokens, block, 0, stream>>>(
+            outputs.data_ptr<scalar_t>(),
+            inputs.data_ptr<scalar_t>(),
+            dst_to_src.data_ptr<int32_t>(),
+            dst_tokens, hidden_size);
+    });
+}
+
+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,
+    std::string format,
+    int64_t persistent,
+    int64_t output_n
+) {
+    // Implementation: loop over experts, call cuinferCustomGemm for each
+    // weights: [num_experts, N, K] with format "TN" means transB
+    // For each expert e with count tokens:
+    //   A = inputs[offset:offset+count, :]        (count × K, row-major)
+    //   B = weights[e, :, :]                      (N × K, needs transB)
+    //   C = output[offset:offset+count, :]        (count × N, row-major)
+    //   GEMM: C = A × B^T  → (count, K) × (K, N) = (count, N)
+
+    auto stream = c10::cuda::getCurrentCUDAStream();
+    int num_experts = weights.size(0);
+    int N = weights.size(1);  // output dim
+    int K = weights.size(2);  // input dim
+
+    // Get token counts on CPU
+    auto counts_cpu = tokens_per_experts.to(torch::kCPU).to(torch::kInt32);
+    int32_t* counts = counts_cpu.data_ptr<int32_t>();
+
+    // Create cuinfer handle
+    cuinferHandle_t handle;
+    cuinferCreate(&handle);
+    cuinferSetStream(handle, stream);
+
+    float alpha = 1.0f, beta = 0.0f;
+
+    int offset = 0;
+    for (int e = 0; e < num_experts; e++) {
+        int M = counts[e];
+        if (M <= 0) continue;
+
+        // A: inputs[offset : offset+M, :]  → M × K
+        // B: weights[e, :, :]               → N × K (transposed: compute A × B^T)
+        // C: output[offset : offset+M, :]   → M × N
+        const void* A_ptr = (const char*)inputs.data_ptr() +
+            (int64_t)offset * K * inputs.element_size();
+        const void* B_ptr = (const char*)weights.data_ptr() +
+            (int64_t)e * N * K * weights.element_size();
+        void* C_ptr = (char*)output.data_ptr() +
+            (int64_t)offset * N * output.element_size();
+
+        cudaDataType_t dtype = (inputs.scalar_type() == torch::kFloat16)
+            ? CUDA_R_16F : CUDA_R_32F;
+
+        // cuinferCustomGemm: row-major convention
+        // We want C = A × B^T
+        // In cuinfer (column-major internally): transa=N, transb=T
+        // M_gemm = M (rows of C), N_gemm = N (cols of C), K_gemm = K
+        cuinferCustomGemm(
+            handle, stream,
+            CUINFER_POINTER_MODE_HOST,
+            CUINFER_OP_TENSOR_OP_N,  // transa = no transpose
+            CUINFER_OP_TENSOR_OP_T,  // transb = transpose (TN format)
+            M, N, K,
+            &alpha,
+            A_ptr, dtype, K, 0,       // lda=K for row-major A
+            B_ptr, dtype, K, 0,       // ldb=K for row-major B (will be transposed)
+            &beta,
+            C_ptr, dtype, N, 0,       // ldc=N for row-major C
+            1,                         // batchCount=1
+            CUDA_R_32F,               // computeType
+            CUDA_R_32F,               // scaleType
+            nullptr, nullptr,          // custom pointers
+            CUINFER_GEMM_DEFAULT);
+
+        offset += M;
+    }
+
+    cuinferDestroy(handle);
+}
+
+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,
+    double scaling_factor
+) {
+    // inputs: [N, topk, H] — expert outputs per token
+    // mul_weight: [N, topk]  — router weights
+    // outputs: [N, H]       — weighted sum
+    auto stream = c10::cuda::getCurrentCUDAStream();
+    int num_tokens = inputs.size(0);
+    int topk = inputs.size(1);
+    int hidden_size = inputs.size(2);
+    int block = std::min(hidden_size, 256);
+
+    // Reshape inputs to [N*topk, H] for the kernel
+    auto input_flat = inputs.reshape({num_tokens * topk, hidden_size});
+
+    if (inputs.scalar_type() == torch::kFloat16) {
+        combine_result_kernel<__half><<<num_tokens, block, 0, stream>>>(
+            reinterpret_cast<__half*>(outputs.data_ptr()),
+            reinterpret_cast<const __half*>(input_flat.data_ptr()),
+            mul_weight.value().data_ptr<float>(),
+            num_tokens, topk, hidden_size);
+    } else {
+        combine_result_kernel<float><<<num_tokens, block, 0, stream>>>(
+            outputs.data_ptr<float>(),
+            input_flat.data_ptr<float>(),
+            mul_weight.value().data_ptr<float>(),
+            num_tokens, topk, hidden_size);
+    }
+}
+
+}}  // namespace ixformer::infer

ex_engine/csrc/moe_tcu_dispatch.cpp

@@ -0,0 +1,191 @@
+// moe_tcu_dispatch.cpp — MoE expert GEMM via torch::mm (walks Gemm_tcu_bi_kernel)
+//
+// Replaces Python for-loop over experts with C++ loop.
+// torch::mm on corex launches Gemm_tcu_bi_kernel::gemm_h_h_tcu_25 (TCU hardware).
+// Probe confirmed: Python loop overhead = 0.892 ms/expert = 7.1 ms for 8 experts.
+// This C++ dispatch eliminates that overhead.
+//
+// No custom GEMM kernel. No ixformer API dependency. Just torch::mm in C++.
+
+#include <torch/extension.h>
+#include <vector>
+
+// ============================================================================
+// Decode path: single token, top_k experts
+// ============================================================================
+// hidden:         (1, K)
+// gate_up_weights: (num_experts, 2*intermediate, K) — pre-loaded expert weights
+// down_weights:    (num_experts, K, intermediate)
+// expert_ids:     (top_k,) int64 — selected expert indices
+// expert_weights: (top_k,) float — gating weights
+//
+// For each expert:
+//   gate_up = hidden @ gate_up_weights[eid].t()   → (1, 2*I)
+//   gate = silu(gate_up[:, :I])
+//   up   = gate_up[:, I:]
+//   act  = gate * up                               → (1, I)
+//   out  = act @ down_weights[eid].t()             → (1, K)
+//   result += weight * out
+
+torch::Tensor moe_decode(
+    torch::Tensor hidden,          // (1, K)
+    torch::Tensor gate_up_weights, // (E, 2*I, K)
+    torch::Tensor down_weights,    // (E, K, I)
+    torch::Tensor expert_ids,      // (top_k,) int64
+    torch::Tensor expert_weights   // (top_k,) float/half
+) {
+    auto top_k = expert_ids.size(0);
+    auto K = hidden.size(1);
+    auto inter2 = gate_up_weights.size(1);
+    auto inter = inter2 / 2;
+
+    auto result = torch::zeros_like(hidden);  // (1, K)
+
+    for (int64_t k = 0; k < top_k; ++k) {
+        auto eid = expert_ids[k].item<int64_t>();
+        auto w = expert_weights[k].item<float>();
+
+        // FC1: gate_up = hidden @ w13[eid]^T → (1, 2*I)
+        auto gate_up = torch::mm(hidden, gate_up_weights[eid].t());
+
+        // SiLU and mul
+        auto gate_slice = gate_up.slice(1, 0, inter); auto gate = gate_slice * torch::sigmoid(gate_slice);
+        auto up = gate_up.slice(1, inter, inter2);
+        auto act = gate * up;  // (1, I)
+
+        // FC2: expert_out = act @ w2[eid]^T → (1, K)
+        auto expert_out = torch::mm(act, down_weights[eid].t());
+
+        // Weighted accumulate
+        result.add_(expert_out, w);
+    }
+
+    return result;
+}
+
+
+// ============================================================================
+// Prefill path: multiple tokens, grouped by expert
+// ============================================================================
+// hidden:          (T, K)
+// gate_up_weights: (E, 2*I, K)
+// down_weights:    (E, K, I)
+// topk_ids:        (T, top_k) int64 — expert indices per token
+// topk_weights:    (T, top_k) float — gating weights per token
+//
+// Strategy: group tokens by expert, batch the GEMM per expert.
+
+torch::Tensor moe_prefill(
+    torch::Tensor hidden,          // (T, K)
+    torch::Tensor gate_up_weights, // (E, 2*I, K)
+    torch::Tensor down_weights,    // (E, K, I)
+    torch::Tensor topk_ids,        // (T, top_k) int64
+    torch::Tensor topk_weights     // (T, top_k) float/half
+) {
+    auto T = hidden.size(0);
+    auto K = hidden.size(1);
+    auto num_experts = gate_up_weights.size(0);
+    auto inter2 = gate_up_weights.size(1);
+    auto inter = inter2 / 2;
+    auto top_k = topk_ids.size(1);
+
+    auto result = torch::zeros({T, K}, hidden.options());
+
+    // Flatten topk_ids to find tokens per expert
+    auto flat_ids = topk_ids.reshape(-1);         // (T*top_k,)
+    auto flat_weights = topk_weights.reshape(-1);  // (T*top_k,)
+
+    // Token index for each (token, k) pair
+    auto token_idx = torch::arange(T, topk_ids.options())
+                         .unsqueeze(1).expand({T, top_k}).reshape(-1);  // (T*top_k,)
+
+    for (int64_t eid = 0; eid < num_experts; ++eid) {
+        // Find which entries in flat_ids match this expert
+        auto mask = flat_ids.eq(eid);
+        auto count = mask.sum().item<int64_t>();
+        if (count == 0) continue;
+
+        // Gather token indices and weights for this expert
+        auto indices = mask.nonzero().squeeze(1);        // (count,)
+        auto tok_indices = token_idx.index_select(0, indices);  // (count,)
+        auto weights = flat_weights.index_select(0, indices);   // (count,)
+
+        // Gather hidden states
+        auto tokens = hidden.index_select(0, tok_indices);  // (count, K)
+
+        // FC1: gate_up = tokens @ w13[eid]^T → (count, 2*I)
+        auto gate_up = torch::mm(tokens, gate_up_weights[eid].t());
+
+        // SiLU and mul
+        auto gate_slice = gate_up.slice(1, 0, inter); auto gate = gate_slice * torch::sigmoid(gate_slice);
+        auto up = gate_up.slice(1, inter, inter2);
+        auto act = gate * up;  // (count, I)
+
+        // FC2: expert_out = act @ w2[eid]^T → (count, K)
+        auto expert_out = torch::mm(act, down_weights[eid].t());
+
+        // Weighted scatter-add
+        auto weighted = expert_out * weights.unsqueeze(1);
+        result.index_add_(0, tok_indices, weighted.to(result.dtype()));
+    }
+
+    return result;
+}
+
+
+// ============================================================================
+// Simple expert GEMM only (no activation, for benchmarking)
+// ============================================================================
+// input:         (total_tokens, K)
+// weights:       (num_experts, N, K)
+// expert_counts: (num_experts,) int64
+// Returns:       (total_tokens, N)
+
+torch::Tensor moe_expert_gemm_tcu(
+    torch::Tensor input,
+    torch::Tensor weights,
+    torch::Tensor expert_counts
+) {
+    auto total_tokens = input.size(0);
+    auto K = input.size(1);
+    auto num_experts = weights.size(0);
+    auto N = weights.size(1);
+
+    auto output = torch::zeros({total_tokens, N}, input.options());
+
+    int64_t offset = 0;
+    for (int64_t e = 0; e < num_experts; ++e) {
+        auto count = expert_counts[e].item<int64_t>();
+        if (count == 0) continue;
+
+        auto tokens = input.slice(0, offset, offset + count);  // (count, K)
+        auto w = weights[e];  // (N, K)
+
+        // torch::mm → Gemm_tcu_bi_kernel on BI-V100
+        auto out_e = torch::mm(tokens, w.t());  // (count, N)
+        output.slice(0, offset, offset + count).copy_(out_e);
+
+        offset += count;
+    }
+
+    return output;
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("moe_decode", &moe_decode,
+          "MoE decode: C++ loop over experts via torch::mm (TCU kernel)",
+          py::arg("hidden"), py::arg("gate_up_weights"),
+          py::arg("down_weights"), py::arg("expert_ids"),
+          py::arg("expert_weights"));
+
+    m.def("moe_prefill", &moe_prefill,
+          "MoE prefill: group-by-expert via torch::mm (TCU kernel)",
+          py::arg("hidden"), py::arg("gate_up_weights"),
+          py::arg("down_weights"), py::arg("topk_ids"),
+          py::arg("topk_weights"));
+
+    m.def("moe_expert_gemm_tcu", &moe_expert_gemm_tcu,
+          "MoE expert GEMM only via torch::mm (TCU kernel, for benchmarking)",
+          py::arg("input"), py::arg("weights"), py::arg("expert_counts"));
+}

ex_engine/deploy_ilu_pipeline.sh

@@ -0,0 +1,188 @@
+#!/usr/bin/env bash
+# deploy_ilu_pipeline.sh — Build + deploy the complete ILU kernel pipeline
+#
+# This replaces ALL Python fallbacks with C++ calls through ixformer::infer.
+# Call from patch_ops.sh after basic vllm patching is done.
+#
+# What this does:
+#   1. Build ix_full_bridge_v2.so (pybind11 bridge to all 14 ixformer functions)
+#   2. Deploy Python dispatch modules (ix_ops_dispatch, corex_gdn, corex_moe, corex_fa2)
+#   3. Deploy upstream xllm ILU kernel wrappers
+#   4. Wire ix_startup_patch to auto-load at vllm import
+#
+# Usage:
+#   bash deploy_ilu_pipeline.sh <VLLM_ROOT>
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+VLLM_ROOT="${1:?Usage: deploy_ilu_pipeline.sh <VLLM_ROOT>}"
+
+echo "============================================"
+echo "[ILU] Starting ILU pipeline deployment"
+echo "[ILU] VLLM_ROOT: ${VLLM_ROOT}"
+echo "[ILU] Script dir: ${SCRIPT_DIR}"
+echo "============================================"
+
+# --- Step 1: Create ex_engine package in vllm ---
+EX_DIR="${VLLM_ROOT}/ex_engine"
+mkdir -p "${EX_DIR}/python"
+cat > "${EX_DIR}/__init__.py" << 'EOF'
+"""ex_engine — Algorithm factor replacement for BI-V100."""
+EOF
+cat > "${EX_DIR}/python/__init__.py" << 'EOF'
+"""ex_engine.python — Python dispatch modules."""
+EOF
+
+# --- Step 2: Try to build ix_full_bridge_v2.so ---
+echo "[ILU] Step 2: Building ix_full_bridge_v2.so..."
+BRIDGE_SO="${SCRIPT_DIR}/prebuilt/ix_full_bridge_v2.so"
+if [[ -f "$BRIDGE_SO" ]]; then
+    echo "[ILU] ✓ Using prebuilt ix_full_bridge_v2.so"
+else
+    if bash "${SCRIPT_DIR}/build_ix_bridge.sh" "${VLLM_ROOT}" 2>&1; then
+        echo "[ILU] ✓ Built ix_full_bridge_v2.so"
+    else
+        echo "[ILU] ⚠ ix_full_bridge_v2.so build failed — will use ixformer Python path"
+    fi
+fi
+
+# Deploy bridge .so
+if [[ -f "$BRIDGE_SO" ]]; then
+    cp "$BRIDGE_SO" "${EX_DIR}/ix_full_bridge_v2.so"
+    cp "$BRIDGE_SO" "${EX_DIR}/python/ix_full_bridge_v2.so"
+    echo "[ILU] ✓ Deployed ix_full_bridge_v2.so"
+fi
+
+# --- Step 3: Deploy Python dispatch modules ---
+echo "[ILU] Step 3: Deploying Python dispatch modules..."
+
+for pyfile in \
+    ix_ops_dispatch.py \
+    corex_gdn.py \
+    corex_moe.py \
+    corex_fa2.py \
+    corex_fa2_dispatch.py \
+    fused_moe_ilu.py \
+    ix_bridge.py \
+    ix_bridge_v2.py \
+    ix_ops.py \
+    patch_vllm_ops.py \
+    ex_loader.py \
+    moe_topk.py \
+    patch_model.py; do
+    src="${SCRIPT_DIR}/python/${pyfile}"
+    if [[ -f "$src" ]]; then
+        cp "$src" "${EX_DIR}/python/${pyfile}"
+        echo "[ILU]   ✓ ${pyfile}"
+    fi
+done
+
+# Also deploy corex_gdn.py and corex_moe.py to vllm models dir for import
+MODELS_DIR="${VLLM_ROOT}/model_executor/models"
+for pyfile in corex_gdn.py corex_moe.py corex_fa2.py; do
+    src="${SCRIPT_DIR}/python/${pyfile}"
+    if [[ -f "$src" ]] && [[ -d "$MODELS_DIR" ]]; then
+        cp "$src" "${MODELS_DIR}/${pyfile}"
+        echo "[ILU]   ✓ ${pyfile} → models/"
+    fi
+done
+
+# --- Step 4: Deploy xllm ILU kernel wrappers ---
+echo "[ILU] Step 4: Deploying xllm ILU kernel sources..."
+ILU_SRC="${SCRIPT_DIR}/xllm_kernels/ilu"
+ILU_UPSTREAM="${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/ilu"
+
+# Copy from upstream if not already in ex_engine
+if [[ -d "$ILU_UPSTREAM" ]] && [[ ! -d "$ILU_SRC" ]]; then
+    mkdir -p "$ILU_SRC"
+    cp "$ILU_UPSTREAM"/*.cpp "$ILU_UPSTREAM"/*.h "$ILU_SRC/" 2>/dev/null || true
+    echo "[ILU]   ✓ Copied from upstream xllm/core/kernels/ilu/"
+fi
+
+if [[ -d "$ILU_SRC" ]]; then
+    mkdir -p "${EX_DIR}/xllm_kernels/ilu"
+    cp "$ILU_SRC"/*.cpp "$ILU_SRC"/*.h "${EX_DIR}/xllm_kernels/ilu/" 2>/dev/null || true
+    echo "[ILU]   ✓ ILU kernel sources deployed"
+fi
+
+# --- Step 5: Deploy upstream kernel sources for reference ---
+echo "[ILU] Step 5: Deploying upstream kernel references..."
+CUDA_SRC="${REPO_ROOT}/upstream_ref/xllm/xllm/core/kernels/cuda"
+if [[ -d "$CUDA_SRC" ]]; then
+    mkdir -p "${EX_DIR}/xllm_kernels/cuda"
+    # Only copy the key files we need
+    for cufile in \
+        activation.cu norm.cu fused_qknorm_rope.cu \
+        reshape_paged_cache.cu block_copy.cu matmul.cpp; do
+        if [[ -f "${CUDA_SRC}/${cufile}" ]]; then
+            cp "${CUDA_SRC}/${cufile}" "${EX_DIR}/xllm_kernels/cuda/"
+        fi
+    done
+    # MoE kernels
+    if [[ -d "${CUDA_SRC}/moe" ]]; then
+        mkdir -p "${EX_DIR}/xllm_kernels/cuda/moe"
+        cp "${CUDA_SRC}/moe"/*.cu "${CUDA_SRC}/moe"/*.cpp \
+           "${EX_DIR}/xllm_kernels/cuda/moe/" 2>/dev/null || true
+    fi
+    # xattention kernels
+    if [[ -d "${CUDA_SRC}/xattention" ]]; then
+        mkdir -p "${EX_DIR}/xllm_kernels/cuda/xattention"
+        cp "${CUDA_SRC}/xattention"/*.cu "${CUDA_SRC}/xattention"/*.cpp \
+           "${CUDA_SRC}/xattention"/*.h \
+           "${EX_DIR}/xllm_kernels/cuda/xattention/" 2>/dev/null || true
+    fi
+    echo "[ILU]   ✓ Upstream CUDA kernel sources deployed"
+fi
+
+# --- Step 6: Deploy ds_vllm libtorch_stable kernels ---
+echo "[ILU] Step 6: Deploying ds_vllm kernel references..."
+DS_SRC="${REPO_ROOT}/upstream_ref/ds_vllm/csrc/libtorch_stable"
+if [[ -d "$DS_SRC" ]]; then
+    mkdir -p "${EX_DIR}/ds_kernels"
+    for cufile in \
+        activation_kernels.cu layernorm_kernels.cu \
+        pos_encoding_kernels.cu cache_kernels.cu; do
+        if [[ -f "${DS_SRC}/${cufile}" ]]; then
+            cp "${DS_SRC}/${cufile}" "${EX_DIR}/ds_kernels/"
+        fi
+    done
+    if [[ -d "${DS_SRC}/moe" ]]; then
+        mkdir -p "${EX_DIR}/ds_kernels/moe"
+        cp "${DS_SRC}/moe/topk_softmax_kernels.cu" \
+           "${DS_SRC}/moe/moe_align_sum_kernels.cu" \
+           "${DS_SRC}/moe/torch_bindings.cpp" \
+           "${EX_DIR}/ds_kernels/moe/" 2>/dev/null || true
+    fi
+    if [[ -d "${DS_SRC}/attention" ]]; then
+        mkdir -p "${EX_DIR}/ds_kernels/attention"
+        cp "${DS_SRC}/attention"/*.cu "${DS_SRC}/attention"/*.cuh \
+           "${EX_DIR}/ds_kernels/attention/" 2>/dev/null || true
+    fi
+    echo "[ILU]   ✓ ds_vllm kernel sources deployed"
+fi
+
+# --- Step 7: Verification ---
+echo "[ILU] Step 7: Verifying deployment..."
+echo "[ILU] ex_engine contents:"
+find "${EX_DIR}" -name "*.py" -o -name "*.so" -o -name "*.cpp" -o -name "*.cu" | sort | head -40
+echo "[ILU] ..."
+COUNT=$(find "${EX_DIR}" -type f | wc -l)
+echo "[ILU] Total files deployed: ${COUNT}"
+
+echo ""
+echo "============================================"
+echo "[ILU] ✓ ILU pipeline deployment complete"
+echo "[ILU] Deployed to: ${EX_DIR}"
+echo "[ILU] "
+echo "[ILU] Runtime dispatch chain:"
+echo "[ILU]   vllm import → ix_startup_patch → patch_vllm_ops"
+echo "[ILU]     → ix_ops_dispatch → ix_full_bridge_v2.so"
+echo "[ILU]       → ixformer::infer::* (C++ kernels)"
+echo "[ILU] "
+echo "[ILU] MoE pipeline:"
+echo "[ILU]   corex_moe.py / fused_moe_ilu.py"
+echo "[ILU]     → topk_softmax → moe_gen_idx → expand → gemm → silu → gemm → combine"
+echo "[ILU]     → ALL through ixformer::infer (no Python expert loop)"
+echo "============================================"

ex_engine/deploy_ix_bridge.sh

@@ -0,0 +1,154 @@
+#!/bin/bash
+# ex_engine/deploy_ix_bridge.sh — Deploy ix_full_bridge + Python ops into vllm
+#
+# Architecture (CCCL build pattern):
+#   CCCL: cmake → compile → install to site-packages
+#   EX:   torch.utils.cpp_extension → compile bridge → deploy to vllm pkg
+#
+# What this does:
+#   1. Find ixformer .so libraries in base image
+#   2. Either use prebuilt ix_full_bridge.so or JIT-compile from source
+#   3. Deploy .so + Python modules into vllm package
+#   4. Verify dlopen chain works
+#
+# Source mapping:
+#   ex_engine/csrc/ix_full_bridge_v2.cpp → pybind11 bridge to ixformer::infer
+#   ex_engine/python/ix_ops.py           → Python API layer
+#   ex_engine/python/patch_vllm_ops.py   → vllm monkey-patches
+#
+# Called from: qwen3_6_scripts/patch_ops.sh
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+VLLM_ROOT="${1:-$(python3 -c 'import vllm; import os; print(os.path.dirname(vllm.__file__))' 2>/dev/null || echo '/usr/local/corex/lib/python3/dist-packages/vllm')}"
+
+echo "[ix_bridge] VLLM_ROOT=${VLLM_ROOT}"
+echo "[ix_bridge] SCRIPT_DIR=${SCRIPT_DIR}"
+
+# =========================================================================
+# Step 1: Deploy prebuilt .so if available
+# =========================================================================
+PREBUILT="${SCRIPT_DIR}/../qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10"
+BRIDGE_SO=""
+
+if [[ -f "${PREBUILT}/ix_full_bridge.so" ]]; then
+    cp "${PREBUILT}/ix_full_bridge.so" "${VLLM_ROOT}/ix_full_bridge.so"
+    BRIDGE_SO="${VLLM_ROOT}/ix_full_bridge.so"
+    echo "[ix_bridge] deployed prebuilt ix_full_bridge.so"
+fi
+
+# Deploy all corex_*.so and xllm_*.so
+if [[ -d "$PREBUILT" ]]; then
+    for so_file in "${PREBUILT}"/*.so; do
+        base=$(basename "$so_file")
+        if [[ "$base" != "ix_full_bridge.so" ]]; then
+            cp "$so_file" "${VLLM_ROOT}/${base}" 2>/dev/null || true
+            echo "[ix_bridge] deployed ${base}"
+        fi
+    done
+fi
+
+# =========================================================================
+# Step 2: Deploy Python integration modules
+# =========================================================================
+# Create ex_engine package in vllm
+EX_PKG="${VLLM_ROOT}/ex_engine"
+mkdir -p "${EX_PKG}"
+
+cat > "${EX_PKG}/__init__.py" << 'PYEOF'
+"""ex_engine — Algorithm factor replacement engine for BI-V100."""
+PYEOF
+
+# Deploy ix_ops.py
+cp "${SCRIPT_DIR}/python/ix_ops.py" "${EX_PKG}/ix_ops.py"
+echo "[ix_bridge] deployed ix_ops.py"
+
+# Deploy patch_vllm_ops.py
+cp "${SCRIPT_DIR}/python/patch_vllm_ops.py" "${EX_PKG}/patch_vllm_ops.py"
+echo "[ix_bridge] deployed patch_vllm_ops.py"
+
+# Also make ix_ops importable from vllm.ex_engine
+# and from the top-level ex_engine path
+SITE_EX="${SCRIPT_DIR}/python"
+if [[ -d "$SITE_EX" ]]; then
+    # Ensure __init__.py exists
+    touch "${SITE_EX}/../__init__.py" 2>/dev/null || true
+fi
+
+# =========================================================================
+# Step 3: Create auto-patch entry point
+# =========================================================================
+# This script is sourced by patch_ops.sh to ensure ix_ops patches
+# are applied at vllm startup
+cat > "${VLLM_ROOT}/ix_startup_patch.py" << 'PYEOF'
+"""
+ix_startup_patch.py — Apply ix_ops patches at vllm startup.
+
+Import this module early in the vllm startup to replace PyTorch fallbacks
+with fused C++ kernels from the base image.
+
+Architecture (CCCL dispatch pattern):
+    import vllm → vllm.__init__ → ix_startup_patch → patch_vllm_ops
+"""
+import logging
+logger = logging.getLogger("ix_startup_patch")
+
+def apply():
+    """Apply all available ix_ops patches."""
+    try:
+        from vllm.ex_engine.patch_vllm_ops import apply_all_patches
+        n = apply_all_patches()
+        if n > 0:
+            logger.info("ix_startup_patch: %d patches applied", n)
+        return n
+    except Exception as e:
+        logger.warning("ix_startup_patch failed: %s", e)
+        return 0
+
+# Auto-apply on import
+_n_patches = apply()
+PYEOF
+echo "[ix_bridge] deployed ix_startup_patch.py"
+
+# =========================================================================
+# Step 4: Deploy bridge C++ source for JIT fallback
+# =========================================================================
+CSRC_DEST="${VLLM_ROOT}/ex_engine/csrc"
+mkdir -p "${CSRC_DEST}"
+for cpp in "${SCRIPT_DIR}/csrc/ix_full_bridge_v2.cpp" \
+           "${SCRIPT_DIR}/csrc/ix_full_bridge.cpp" \
+           "${SCRIPT_DIR}/csrc/ix_moe_bridge.cpp"; do
+    if [[ -f "$cpp" ]]; then
+        cp "$cpp" "${CSRC_DEST}/"
+        echo "[ix_bridge] deployed $(basename $cpp) for JIT fallback"
+    fi
+done
+
+# =========================================================================
+# Step 5: Verify deployment
+# =========================================================================
+echo ""
+echo "[ix_bridge] === Deployment Summary ==="
+echo "[ix_bridge] Bridge .so: ${BRIDGE_SO:-'(JIT compile at runtime)'}"
+echo "[ix_bridge] Python ops: ${EX_PKG}/ix_ops.py"
+echo "[ix_bridge] vllm patches: ${EX_PKG}/patch_vllm_ops.py"
+echo "[ix_bridge] Startup hook: ${VLLM_ROOT}/ix_startup_patch.py"
+
+# Quick Python import test
+python3 -c "
+import sys
+sys.path.insert(0, '${VLLM_ROOT}')
+try:
+    from vllm.ex_engine import ix_ops
+    print('[ix_bridge] ✓ ix_ops importable')
+except Exception as e:
+    print(f'[ix_bridge] ✗ ix_ops import failed: {e}')
+try:
+    from vllm.ex_engine import patch_vllm_ops
+    print('[ix_bridge] ✓ patch_vllm_ops importable')
+except Exception as e:
+    print(f'[ix_bridge] ✗ patch_vllm_ops import failed: {e}')
+" 2>&1 || true
+
+echo "[ix_bridge] === Done ==="

ex_engine/fla_kernels/gated_delta_rule/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+from .chunk import chunk_gated_delta_rule, chunk_gdn
+from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gdn
+from .naive import naive_chunk_gated_delta_rule, naive_recurrent_gated_delta_rule
+
+__all__ = [
+    "chunk_gated_delta_rule", "chunk_gdn",
+    "fused_recurrent_gated_delta_rule", "fused_recurrent_gdn",
+    "naive_chunk_gated_delta_rule",
+    "naive_recurrent_gated_delta_rule",
+]

ex_engine/fla_kernels/gated_delta_rule/chunk.py

@@ -0,0 +1,591 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import warnings
+
+import torch
+
+from fla.modules.l2norm import l2norm_bwd, l2norm_fwd
+from fla.ops.backends import dispatch
+from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
+from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o
+from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd
+from fla.ops.cp import FLACPContext
+from fla.ops.cp.chunk_delta_h import (
+    chunk_gated_delta_rule_bwd_dhu_pre_process,
+    chunk_gated_delta_rule_fwd_h_pre_process,
+    compress_h0,
+    expand_h0,
+)
+from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra
+from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum
+from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd
+from fla.ops.utils import chunk_local_cumsum
+from fla.ops.utils.constant import RCP_LN2
+from fla.ops.utils.index import prepare_chunk_indices
+from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
+
+
+def chunk_gated_delta_rule_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    state_v_first: bool = False,
+    cu_seqlens: torch.LongTensor | None = None,
+    cp_context: FLACPContext | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    use_gate_in_kernel: bool = False,
+    A_log: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    chunk_size: int = 64,
+):
+    g_input = g if use_gate_in_kernel else None
+    if use_gate_in_kernel:
+        g = gdn_gate_chunk_cumsum(
+            g=g,
+            A_log=A_log,
+            chunk_size=chunk_size,
+            scale=RCP_LN2,
+            dt_bias=dt_bias,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+        )
+    else:
+        g = chunk_local_cumsum(
+            g,
+            chunk_size=chunk_size,
+            scale=RCP_LN2,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+        )
+    # obtain WY representation. u is actually the new v.
+    # fused kkt + solve_tril + recompute_w_u
+    w, u, A = chunk_gated_delta_rule_fwd_intra(
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+    )
+
+    if cp_context is not None:
+        initial_state = chunk_gated_delta_rule_fwd_h_pre_process(
+            k=k,
+            w=w,
+            u=u,
+            g=g,
+            cu_seqlens=cu_seqlens,
+            initial_state=initial_state,
+            context=cp_context,
+            state_v_first=state_v_first,
+            chunk_size=chunk_size,
+        )
+
+    h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
+        k=k,
+        w=w,
+        u=u,
+        g=g,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        state_v_first=state_v_first,
+        chunk_size=chunk_size,
+    )
+
+    if cp_context is not None:
+        initial_state = compress_h0(initial_state, context=cp_context)
+
+    o = chunk_fwd_o(
+        q=q,
+        k=k,
+        v=v_new,
+        h=h,
+        g=g,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        state_v_first=state_v_first,
+        chunk_size=chunk_size,
+    )
+    return g, o, A, final_state, initial_state, g_input
+
+
+def chunk_gated_delta_rule_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    do: torch.Tensor,
+    dht: torch.Tensor,
+    state_v_first: bool = False,
+    cu_seqlens: torch.LongTensor | None = None,
+    cp_context: FLACPContext | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    use_gate_in_kernel: bool = False,
+    g_input: torch.Tensor | None = None,
+    A_log: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    chunk_size: int = 64,
+):
+    w, u = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        g=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+
+    if cp_context is not None:
+        initial_state = expand_h0(initial_state, context=cp_context)
+
+    h, v_new, _ = chunk_gated_delta_rule_fwd_h(
+        k=k,
+        w=w,
+        u=u,
+        g=g,
+        initial_state=initial_state,
+        output_final_state=False,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        state_v_first=state_v_first,
+        chunk_size=chunk_size,
+    )
+    dv = chunk_bwd_dv_local(
+        q=q,
+        k=k,
+        g=g,
+        do=do,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_size=chunk_size,
+    )
+
+    if cp_context is not None:
+        # initial_state is None in the CP mode
+        # We only need to compute dht of current rank and pass it to the backward kernel
+        dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process(
+            q=q,
+            k=k,
+            w=w,
+            do=do,
+            dv=dv,
+            g=g,
+            scale=scale,
+            cu_seqlens=cu_seqlens,
+            dht=dht,
+            initial_state=initial_state,
+            context=cp_context,
+            state_v_first=state_v_first,
+            chunk_size=chunk_size,
+        )
+
+    dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(
+        q=q,
+        k=k,
+        w=w,
+        g=g,
+        h0=initial_state,
+        dht=dht,
+        do=do,
+        dv=dv,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        state_v_first=state_v_first,
+        chunk_size=chunk_size,
+    )
+    dq, dk, dw, dg = chunk_bwd_dqkwg(
+        q=q,
+        k=k,
+        v=v_new,
+        w=w,
+        g=g,
+        h=h,
+        dv=dv,
+        do=do,
+        dh=dh,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        state_v_first=state_v_first,
+        chunk_size=chunk_size,
+    )
+    dk2, dv, db, dg2 = prepare_wy_repr_bwd(
+        k=k,
+        v=v,
+        beta=beta,
+        g=g,
+        A=A,
+        dw=dw,
+        du=dv,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    dk.add_(dk2)
+    dg.add_(dg2)
+    dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices)
+    dA_log, ddt_bias = None, None
+    if use_gate_in_kernel:
+        dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg)
+    return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias
+
+
+class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        scale: float,
+        initial_state: torch.Tensor,
+        output_final_state: bool,
+        state_v_first: bool = False,
+        cu_seqlens: torch.LongTensor | None = None,
+        cu_seqlens_cpu: torch.LongTensor | None = None,
+        use_qk_l2norm_in_kernel: bool = False,
+        use_gate_in_kernel: bool = False,
+        A_log: torch.Tensor | None = None,
+        dt_bias: torch.Tensor | None = None,
+        use_beta_sigmoid_in_kernel: bool = False,
+        allow_neg_eigval: bool = False,
+        cp_context: FLACPContext | None = None,
+        chunk_size: int = 64,
+    ):
+        q_rstd, k_rstd = None, None
+        if use_qk_l2norm_in_kernel:
+            q, q_rstd = l2norm_fwd(q)
+            k, k_rstd = l2norm_fwd(k)
+
+        beta_raw = beta
+        if use_beta_sigmoid_in_kernel:
+            beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0)
+
+        chunk_indices = None
+        if cu_seqlens is not None:
+            chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu)
+        g, o, A, final_state, initial_state, g_input = chunk_gated_delta_rule_fwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            scale=scale,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            cu_seqlens=cu_seqlens,
+            cp_context=cp_context,
+            chunk_indices=chunk_indices,
+            state_v_first=state_v_first,
+            use_gate_in_kernel=use_gate_in_kernel,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            chunk_size=chunk_size,
+        )
+        ctx.save_for_backward(
+            q,
+            q_rstd,
+            k,
+            k_rstd,
+            v,
+            g,
+            beta_raw,
+            beta,
+            A,
+            initial_state,
+            cu_seqlens,
+            chunk_indices,
+            g_input,
+            A_log,
+            dt_bias,
+        )
+        ctx.scale = scale
+        ctx.chunk_size = chunk_size
+        ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
+        ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel
+        ctx.allow_neg_eigval = allow_neg_eigval
+        ctx.cp_context = cp_context
+        ctx.state_v_first = state_v_first
+        ctx.use_gate_in_kernel = use_gate_in_kernel
+        return o.to(q.dtype), final_state
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(
+        ctx,
+        do: torch.Tensor,
+        dht: torch.Tensor,
+    ):
+        (
+            q,
+            q_rstd,
+            k,
+            k_rstd,
+            v,
+            g,
+            beta_raw,
+            beta,
+            A,
+            initial_state,
+            cu_seqlens,
+            chunk_indices,
+            g_input,
+            A_log,
+            dt_bias,
+        ) = ctx.saved_tensors
+        dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            A=A,
+            scale=ctx.scale,
+            initial_state=initial_state,
+            do=do,
+            dht=dht,
+            cu_seqlens=cu_seqlens,
+            cp_context=ctx.cp_context,
+            chunk_indices=chunk_indices,
+            state_v_first=ctx.state_v_first,
+            use_gate_in_kernel=ctx.use_gate_in_kernel,
+            g_input=g_input,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            chunk_size=ctx.chunk_size,
+        )
+        if ctx.use_qk_l2norm_in_kernel:
+            dq = l2norm_bwd(q, q_rstd, dq)
+            dk = l2norm_bwd(k, k_rstd, dk)
+        if ctx.use_beta_sigmoid_in_kernel:
+            db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0)
+        return (
+            dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta_raw),
+            None, dh0, None, None, None, None, None, None, dA_log, ddt_bias,
+            None, None, None, None,
+        )
+
+
+@dispatch('gated_delta_rule')
+@torch.compiler.disable
+def chunk_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    use_beta_sigmoid_in_kernel: bool = False,
+    allow_neg_eigval: bool = False,
+    state_v_first: bool = False,
+    cu_seqlens: torch.LongTensor | None = None,
+    cu_seqlens_cpu: torch.LongTensor | None = None,
+    cp_context: FLACPContext | None = None,
+    **kwargs,
+):
+    r"""
+    Args:
+        q (torch.Tensor):
+            queries of shape `[B, T, H, K]`.
+        k (torch.Tensor):
+            keys of shape `[B, T, H, K]`.
+        v (torch.Tensor):
+            values of shape `[B, T, HV, V]`.
+            GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.
+        g (torch.Tensor):
+            (forget) gating tensor of shape `[B, T, HV]`.
+            When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay).
+            When `use_gate_in_kernel=True`, `g` is the raw input before gate activation;
+            the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally.
+        beta (torch.Tensor):
+            betas of shape `[B, T, HV]`.
+        scale (Optional[float]):
+            Scale factor for the RetNet attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `[N, HV, K, V]` for `N` input sequences.
+            For equal-length input sequences, `N` equals the batch size `B`.
+            Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
+        use_qk_l2norm_in_kernel (bool):
+            Whether to apply L2norm to the q/k tensor internally. Default: `False`.
+        use_gate_in_kernel (bool):
+            Whether to compute the log-space GDN decay internally.
+            When `True`, the passed `g` is the raw input, and `A_log` must be provided.
+            The kernel fuses gate activation + chunk cumsum in a single pass.
+            Default: `False`.
+        A_log (Optional[torch.Tensor]):
+            Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.
+        dt_bias (Optional[torch.Tensor]):
+            Bias added to `g` before activation, of shape `[HV]`.
+            Only used when `use_gate_in_kernel=True`.
+        use_beta_sigmoid_in_kernel (bool):
+            Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel.
+            - If `True`, the passed `beta` acts as the raw beta logits.
+            - If `False`, `beta` is expected to already be in post-sigmoid space.
+            Default: `False`.
+        allow_neg_eigval (bool):
+            Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.
+            Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case
+            the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.
+        state_v_first (Optional[bool]):
+            Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.
+        cu_seqlens (torch.LongTensor):
+            Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
+            consistent with the FlashAttention API.
+        cp_context (Optional[FLACPContext]):
+            Context parallel context for distributed training across multiple devices.
+            When provided, `initial_state` and `output_final_state` are not supported,
+            and `cu_seqlens` will be overridden by the context. Default: `None`.
+
+    Returns:
+        o (torch.Tensor):
+            Outputs of shape `[B, T, HV, V]`.
+        final_state (torch.Tensor):
+            Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
+
+    Examples::
+        >>> import torch
+        >>> import torch.nn.functional as F
+        >>> from einops import rearrange
+        >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+        # inputs with equal lengths
+        >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
+        >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
+        >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
+        >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')
+        >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid()
+        >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda'))
+        >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')
+        >>> o, ht = chunk_gated_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True
+        )
+        # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
+        >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
+        # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
+        >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
+        >>> o, ht = chunk_gated_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True,
+            cu_seqlens=cu_seqlens
+        )
+    """
+    if 'transpose_state_layout' in kwargs:
+        if state_v_first:
+            raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.")
+        warnings.warn(
+            "`transpose_state_layout` is deprecated and renamed to `state_v_first`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        state_v_first = kwargs.pop('transpose_state_layout')
+
+    # Validate head dimensions
+    if q.shape[2] != k.shape[2]:
+        raise ValueError(
+            f"q and k must have the same number of heads, "
+            f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}"
+        )
+    H, HV = q.shape[2], v.shape[2]
+    if HV % H != 0:
+        raise ValueError(
+            f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by "
+            f"num_heads (H={H}), but got HV % H = {HV % H}"
+        )
+
+    if 'head_first' in kwargs:
+        raise DeprecationWarning(
+            "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
+        )
+
+    chunk_size = kwargs.pop('chunk_size', 64)
+    if chunk_size not in (16, 32, 64):
+        raise ValueError(f"`chunk_size` must be 16, 32, or 64 for Gated Delta Rule, got {chunk_size}.")
+
+    if cp_context is not None:
+        assert initial_state is None, "Initial state is not supported for CP"
+        assert output_final_state is False, "Output final state is not supported for CP"
+        assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP"
+        cu_seqlens = cp_context.cu_seqlens
+        if cp_context.cu_seqlens_cpu is not None:
+            cu_seqlens_cpu = cp_context.cu_seqlens_cpu
+
+    if cu_seqlens is not None:
+        if q.shape[0] != 1:
+            raise ValueError(
+                f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
+                f"Please flatten variable-length inputs before processing.",
+            )
+        if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
+            raise ValueError(
+                f"The number of initial states is expected to be equal to the number of input sequences, "
+                f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.",
+            )
+    use_gate_in_kernel = kwargs.get('use_gate_in_kernel', False)
+    A_log = kwargs.get('A_log')
+    dt_bias = kwargs.get('dt_bias')
+    if use_gate_in_kernel:
+        assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True."
+    if allow_neg_eigval and not use_beta_sigmoid_in_kernel:
+        raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.")
+
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    o, final_state = ChunkGatedDeltaRuleFunction.apply(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        scale,
+        initial_state,
+        output_final_state,
+        state_v_first,
+        cu_seqlens,
+        cu_seqlens_cpu,
+        use_qk_l2norm_in_kernel,
+        use_gate_in_kernel,
+        A_log,
+        dt_bias,
+        use_beta_sigmoid_in_kernel,
+        allow_neg_eigval,
+        cp_context,
+        chunk_size,
+    )
+    return o, final_state
+
+
+chunk_gdn = chunk_gated_delta_rule

ex_engine/fla_kernels/gated_delta_rule/chunk_fwd.py

@@ -0,0 +1,428 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import torch
+import triton
+import triton.language as tl
+
+from fla.ops.backends import dispatch
+from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
+from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd
+from fla.ops.utils import prepare_chunk_indices, solve_tril
+from fla.ops.utils.cache import fla_cache_autotune
+from fla.ops.utils.op import exp2
+from fla.utils import IS_INTEL, IS_TF32_SUPPORTED, autotune_cache_kwargs
+
+if IS_TF32_SUPPORTED:
+    SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32')
+else:
+    SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee')
+
+
+@triton.heuristics({
+    'USE_G': lambda args: args['g'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@fla_cache_autotune(
+    configs=[
+        triton.Config({'BK': BK}, num_warps=num_warps)
+        for BK in [32, 64]
+        for num_warps in [1, 2, 4]
+    ],
+    key=['H', 'HV', 'K', 'BC'],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=['T'])
+def chunk_gated_delta_rule_fwd_kkt_solve_kernel(
+    k,
+    g,
+    beta,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    USE_G: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    """
+    Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass.
+
+    This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel,
+    avoiding the HBM round-trip for the intermediate A matrix.
+
+    Steps:
+    1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers
+    2. Apply gate and beta scaling
+    3. Forward substitution on diagonal blocks
+    4. Block merge to get full (I+A)^{-1}
+    5. Write result to A (output)
+    """
+    i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    i_b, i_h = i_bh // HV, i_bh % HV
+
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    if i_t * BT >= T:
+        return
+
+    i_tc0 = i_t * BT
+    i_tc1 = i_t * BT + BC
+    i_tc2 = i_t * BT + 2 * BC
+    i_tc3 = i_t * BT + 3 * BC
+
+    k += (bos * H + i_h // (HV // H)) * K
+    A += (bos * HV + i_h) * BT
+
+    o_i = tl.arange(0, BC)
+    m_tc0 = (i_tc0 + o_i) < T
+    m_tc1 = (i_tc1 + o_i) < T
+    m_tc2 = (i_tc2 + o_i) < T
+    m_tc3 = (i_tc3 + o_i) < T
+
+    # load beta for each sub-chunk
+    p_b0 = beta + bos * HV + i_h + (i_tc0 + o_i) * HV
+    p_b1 = beta + bos * HV + i_h + (i_tc1 + o_i) * HV
+    p_b2 = beta + bos * HV + i_h + (i_tc2 + o_i) * HV
+    p_b3 = beta + bos * HV + i_h + (i_tc3 + o_i) * HV
+    b_b0 = tl.load(p_b0, mask=m_tc0, other=0.0).to(tl.float32)
+    b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32)
+    b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32)
+    b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32)
+
+    # load gate if used
+    if USE_G:
+        p_g0 = g + bos * HV + i_h + (i_tc0 + o_i) * HV
+        p_g1 = g + bos * HV + i_h + (i_tc1 + o_i) * HV
+        p_g2 = g + bos * HV + i_h + (i_tc2 + o_i) * HV
+        p_g3 = g + bos * HV + i_h + (i_tc3 + o_i) * HV
+
+        b_g0 = tl.load(p_g0, mask=m_tc0, other=0.0).to(tl.float32)
+        b_g1 = tl.load(p_g1, mask=m_tc1, other=0.0).to(tl.float32)
+        b_g2 = tl.load(p_g2, mask=m_tc2, other=0.0).to(tl.float32)
+        b_g3 = tl.load(p_g3, mask=m_tc3, other=0.0).to(tl.float32)
+
+    ############################################################################
+    # Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T
+    ############################################################################
+
+    # 4 diagonal blocks
+    b_A00 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A11 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A22 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A33 = tl.zeros([BC, BC], dtype=tl.float32)
+
+    # 6 off-diagonal blocks
+    b_A10 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A20 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A21 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A30 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A31 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A32 = tl.zeros([BC, BC], dtype=tl.float32)
+
+    for i_k in range(tl.cdiv(K, BK)):
+        o_k = i_k * BK + tl.arange(0, BK)
+        p_k0 = k + (i_tc0 + o_i)[:, None] * (H*K) + o_k[None, :]
+        b_k0 = tl.load(p_k0, mask=m_tc0[:, None] & (o_k[None, :] < K), other=0.0)
+        # diagonal block 0
+        b_A00 += tl.dot(b_k0, tl.trans(b_k0))
+
+        if i_tc1 < T:
+            p_k1 = k + (i_tc1 + o_i)[:, None] * (H*K) + o_k[None, :]
+            b_k1 = tl.load(p_k1, mask=m_tc1[:, None] & (o_k[None, :] < K), other=0.0)
+            # diagonal block 1
+            b_A11 += tl.dot(b_k1, tl.trans(b_k1))
+            # off-diagonal (1,0)
+            b_A10 += tl.dot(b_k1, tl.trans(b_k0))
+
+            if i_tc2 < T:
+                p_k2 = k + (i_tc2 + o_i)[:, None] * (H*K) + o_k[None, :]
+                b_k2 = tl.load(p_k2, mask=m_tc2[:, None] & (o_k[None, :] < K), other=0.0)
+                # diagonal block 2
+                b_A22 += tl.dot(b_k2, tl.trans(b_k2))
+                # off-diagonal (2,0), (2,1)
+                b_A20 += tl.dot(b_k2, tl.trans(b_k0))
+                b_A21 += tl.dot(b_k2, tl.trans(b_k1))
+
+                if i_tc3 < T:
+                    p_k3 = k + (i_tc3 + o_i)[:, None] * (H*K) + o_k[None, :]
+                    b_k3 = tl.load(p_k3, mask=m_tc3[:, None] & (o_k[None, :] < K), other=0.0)
+                    # diagonal block 3
+                    b_A33 += tl.dot(b_k3, tl.trans(b_k3))
+                    # off-diagonal (3,0), (3,1), (3,2)
+                    b_A30 += tl.dot(b_k3, tl.trans(b_k0))
+                    b_A31 += tl.dot(b_k3, tl.trans(b_k1))
+                    b_A32 += tl.dot(b_k3, tl.trans(b_k2))
+
+    ############################################################################
+    # Step 2: apply gate and beta scaling
+    ############################################################################
+
+    # apply gate, beta scaling, and masking
+    # m_d: strictly lower triangular mask for diagonal blocks
+    # m_tc: boundary mask to prevent NaN from 0 * inf (IEEE 754) when
+    #   out-of-bounds g loads as 0 via boundary_check and exp2(0 - g_inbounds) overflows
+    m_d = o_i[:, None] > o_i[None, :]
+    m_I = o_i[:, None] == o_i[None, :]
+
+    if USE_G:
+        b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp2(b_g0[:, None] - b_g0[None, :]), 0.)
+        b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp2(b_g1[:, None] - b_g1[None, :]), 0.)
+        b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp2(b_g2[:, None] - b_g2[None, :]), 0.)
+        b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp2(b_g3[:, None] - b_g3[None, :]), 0.)
+
+        b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp2(b_g1[:, None] - b_g0[None, :]), 0.)
+        b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp2(b_g2[:, None] - b_g0[None, :]), 0.)
+        b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp2(b_g2[:, None] - b_g1[None, :]), 0.)
+        b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp2(b_g3[:, None] - b_g0[None, :]), 0.)
+        b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp2(b_g3[:, None] - b_g1[None, :]), 0.)
+        b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp2(b_g3[:, None] - b_g2[None, :]), 0.)
+    else:
+        b_A00 = tl.where(m_d, b_A00, 0.)
+        b_A11 = tl.where(m_d, b_A11, 0.)
+        b_A22 = tl.where(m_d, b_A22, 0.)
+        b_A33 = tl.where(m_d, b_A33, 0.)
+
+    # diagonal blocks: scaled by beta
+    b_A00 = b_A00 * b_b0[:, None]
+    b_A11 = b_A11 * b_b1[:, None]
+    b_A22 = b_A22 * b_b2[:, None]
+    b_A33 = b_A33 * b_b3[:, None]
+
+    # off-diagonal blocks: full block, scaled by beta
+    b_A10 = b_A10 * b_b1[:, None]
+    b_A20 = b_A20 * b_b2[:, None]
+    b_A21 = b_A21 * b_b2[:, None]
+    b_A30 = b_A30 * b_b3[:, None]
+    b_A31 = b_A31 * b_b3[:, None]
+    b_A32 = b_A32 * b_b3[:, None]
+
+    ############################################################################
+    # Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1}
+    #
+    # Same algorithm as solve_tril, but rows are extracted from in-register
+    # [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of
+    # tl.load from HBM.
+    ############################################################################
+
+    b_Ai00 = -b_A00
+    b_Ai11 = -b_A11
+    b_Ai22 = -b_A22
+    b_Ai33 = -b_A33
+
+    for i in range(2, min(BC, T - i_tc0)):
+        b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.), 0)
+        b_a00 = tl.where(o_i < i, b_a00, 0.)
+        b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0)
+        b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
+    for i in range(2, min(BC, T - i_tc1)):
+        b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.), 0)
+        b_a11 = tl.where(o_i < i, b_a11, 0.)
+        b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0)
+        b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11)
+    for i in range(2, min(BC, T - i_tc2)):
+        b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.), 0)
+        b_a22 = tl.where(o_i < i, b_a22, 0.)
+        b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0)
+        b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22)
+    for i in range(2, min(BC, T - i_tc3)):
+        b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.), 0)
+        b_a33 = tl.where(o_i < i, b_a33, 0.)
+        b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0)
+        b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33)
+
+    b_Ai00 += m_I
+    b_Ai11 += m_I
+    b_Ai22 += m_I
+    b_Ai33 += m_I
+
+    ############################################################################
+    # Step 4: block merge -> full (I + A)^{-1}
+    ############################################################################
+
+    b_Ai10 = -tl.dot(
+        tl.dot(b_Ai11, b_A10, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        b_Ai00,
+        input_precision=SOLVE_TRIL_DOT_PRECISION
+    )
+    b_Ai21 = -tl.dot(
+        tl.dot(b_Ai22, b_A21, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        b_Ai11,
+        input_precision=SOLVE_TRIL_DOT_PRECISION
+    )
+    b_Ai32 = -tl.dot(
+        tl.dot(b_Ai33, b_A32, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        b_Ai22,
+        input_precision=SOLVE_TRIL_DOT_PRECISION
+    )
+
+    b_Ai20 = -tl.dot(
+        b_Ai22,
+        tl.dot(b_A20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
+        tl.dot(b_A21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+    b_Ai31 = -tl.dot(
+        b_Ai33,
+        tl.dot(b_A31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) +
+        tl.dot(b_A32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+    b_Ai30 = -tl.dot(
+        b_Ai33,
+        tl.dot(b_A30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
+        tl.dot(b_A31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) +
+        tl.dot(b_A32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
+        input_precision=SOLVE_TRIL_DOT_PRECISION,
+    )
+
+    ############################################################################
+    # Step 5: store full (I + A)^{-1} to output A
+    ############################################################################
+
+    p_A00 = A + (i_tc0 + o_i)[:, None] * (HV*BT) + o_i[None, :]
+    p_A10 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + o_i[None, :]
+    p_A11 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
+    p_A20 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + o_i[None, :]
+    p_A21 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
+    p_A22 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :]
+    p_A30 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + o_i[None, :]
+    p_A31 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
+    p_A32 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :]
+    p_A33 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (3*BC + o_i)[None, :]
+
+    m_A0 = m_tc0[:, None] & (o_i[None, :] < BT)
+    m_A1 = m_tc1[:, None] & (o_i[None, :] < BT)
+    m_A2 = m_tc2[:, None] & (o_i[None, :] < BT)
+    m_A3 = m_tc3[:, None] & (o_i[None, :] < BT)
+    m_A11 = m_tc1[:, None] & ((BC + o_i)[None, :] < BT)
+    m_A21 = m_tc2[:, None] & ((BC + o_i)[None, :] < BT)
+    m_A22 = m_tc2[:, None] & ((2*BC + o_i)[None, :] < BT)
+    m_A31 = m_tc3[:, None] & ((BC + o_i)[None, :] < BT)
+    m_A32 = m_tc3[:, None] & ((2*BC + o_i)[None, :] < BT)
+    m_A33 = m_tc3[:, None] & ((3*BC + o_i)[None, :] < BT)
+
+    tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), mask=m_A0)
+    tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), mask=m_A1)
+    tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), mask=m_A11)
+    tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), mask=m_A2)
+    tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), mask=m_A21)
+    tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), mask=m_A22)
+    tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), mask=m_A3)
+    tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), mask=m_A31)
+    tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), mask=m_A32)
+    tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), mask=m_A33)
+
+
+@dispatch('gated_delta_rule')
+def chunk_gated_delta_rule_fwd_intra(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor | None = None,
+    beta: torch.Tensor | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    chunk_indices: torch.LongTensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    r"""
+    GDN intra-chunk forward: fused or unfused kkt + solve_tril + recompute_w_u.
+
+    For ``chunk_size == 64``, this uses the fused kkt + solve_tril path. For
+    other supported chunk sizes, it computes the mathematically equivalent
+    representation with ``chunk_scaled_dot_kkt_fwd`` followed by ``solve_tril``.
+
+    Args:
+        k (torch.Tensor):
+            The key tensor of shape `[B, T, H, K]`.
+        v (torch.Tensor):
+            The value tensor of shape `[B, T, HV, V]`.
+        g (torch.Tensor):
+            The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: `None`.
+        beta (torch.Tensor):
+            The beta tensor of shape `[B, T, HV]`.
+        cu_seqlens (torch.LongTensor):
+            The cumulative sequence lengths. Default: `None`.
+        chunk_size (int):
+            The chunk size. Default: 64.
+        chunk_indices (torch.LongTensor):
+            Precomputed chunk indices. Default: `None`.
+
+    Returns:
+        w (torch.Tensor): shape `[B, T, HV, K]`
+        u (torch.Tensor): shape `[B, T, HV, V]`
+        A (torch.Tensor): shape `[B, T, HV, BT]`, the solved (I+A)^{-1} matrix
+    """
+    if chunk_size not in (16, 32, 64):
+        raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.")
+
+    B, T, H, K, HV = *k.shape, beta.shape[2]
+    BT = chunk_size
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+
+    # The fused kernel keeps ten [BC, BC] fp32 accumulators live across the K loop.
+    # That fits NVIDIA's register file but spills on Intel GPUs, where the unfused
+    # two-kernel path measures 2.3-3.0x faster despite the extra HBM round-trip.
+    if BT == 64 and not IS_INTEL:
+        # Step 1: fused kkt + solve_tril
+        BC = 16
+        NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+        A = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype)
+        chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * HV)](
+            k=k,
+            g=g,
+            beta=beta,
+            A=A,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            T=T,
+            H=H,
+            HV=HV,
+            K=K,
+            BT=BT,
+            BC=BC,
+        )
+    else:
+        # Step 1: mathematically equivalent unfused kkt + solve_tril
+        A = chunk_scaled_dot_kkt_fwd(
+            k=k,
+            g=g,
+            beta=beta,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            chunk_size=BT,
+            output_dtype=torch.float32,
+        )
+        A = solve_tril(
+            A=A,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            output_dtype=k.dtype,
+        )
+
+    # Step 2: recompute_w_u
+    w, u = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        g=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    return w, u, A

ex_engine/fla_kernels/gated_delta_rule/fused_recurrent.py

@@ -0,0 +1,478 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import warnings
+
+import torch
+import triton
+import triton.language as tl
+
+from fla.ops.utils.op import exp
+from fla.ops.utils.softplus import softplus
+from fla.utils import input_guard
+
+
+@triton.heuristics({
+    'USE_G': lambda args: args['g'] is not None,
+    'USE_GK': lambda args: args['gk'] is not None,
+    'USE_GV': lambda args: args['gv'] is not None,
+    'USE_INITIAL_STATE': lambda args: args['h0'] is not None,
+    'STORE_FINAL_STATE': lambda args: args['ht'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+    'USE_GATE_IN_KERNEL': lambda args: args['A_log'] is not None,
+    'HAS_DT_BIAS': lambda args: args['dt_bias'] is not None,
+})
+@triton.jit(do_not_specialize=['T'])
+def fused_recurrent_gated_delta_rule_fwd_kernel(
+    q,
+    k,
+    v,
+    g,
+    gk,
+    gv,
+    beta,
+    A_log,
+    dt_bias,
+    o,
+    h0,
+    ht,
+    cu_seqlens,
+    scale,
+    T,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_GK: tl.constexpr,
+    USE_GV: tl.constexpr,
+    USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+    IS_BETA_HEADWISE: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    STORE_FINAL_STATE: tl.constexpr,
+    STATE_V_FIRST: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_GATE_IN_KERNEL: tl.constexpr,
+    HAS_DT_BIAS: tl.constexpr,
+    APPLY_BETA_SIGMOID: tl.constexpr,
+    ALLOW_NEG_EIGVAL: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    NV = tl.cdiv(V, BV)
+    i_v, i_nh = pid % NV, (pid // NV).to(tl.int64)
+    i_n, i_hv = i_nh // HV, i_nh % HV
+    i_h = i_hv // (HV // H)
+
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
+        T = eos - bos
+    else:
+        bos, eos = i_n * T, i_n * T + T
+    o_k = tl.arange(0, BK)
+    o_v = i_v * BV + tl.arange(0, BV)
+
+    p_q = q + (bos * H + i_h) * K + o_k
+    p_k = k + (bos * H + i_h) * K + o_k
+    p_v = v + (bos * HV + i_hv) * V + o_v
+    if USE_G:
+        p_g = g + bos * HV + i_hv
+    if USE_GK:
+        p_gk = gk + (bos * HV + i_hv) * K + o_k
+    if USE_GV:
+        p_gv = gv + (bos * HV + i_hv) * V + o_v
+    if IS_BETA_HEADWISE:
+        p_beta = beta + bos * HV + i_hv
+    else:
+        p_beta = beta + (bos * HV + i_hv) * V + o_v
+
+    p_o = o + (bos * HV + i_hv) * V + o_v
+
+    mask_k = o_k < K
+    mask_v = o_v < V
+    if STATE_V_FIRST:
+        mask_h = mask_v[:, None] & mask_k[None, :]
+    else:
+        mask_h = mask_k[:, None] & mask_v[None, :]
+
+    if STATE_V_FIRST:
+        b_h = tl.zeros([BV, BK], dtype=tl.float32)
+    else:
+        b_h = tl.zeros([BK, BV], dtype=tl.float32)
+    if USE_INITIAL_STATE:
+        if STATE_V_FIRST:
+            p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :]
+        else:
+            p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :]
+        b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
+
+    for _ in tl.range(0, T):
+        b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
+        b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
+        b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
+        if USE_QK_L2NORM_IN_KERNEL:
+            b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
+            b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
+        b_q = b_q * scale
+        if IS_BETA_HEADWISE:
+            b_beta = tl.load(p_beta).to(tl.float32)
+        else:
+            b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
+        if APPLY_BETA_SIGMOID:
+            b_beta = tl.sigmoid(b_beta)
+            if ALLOW_NEG_EIGVAL:
+                b_beta = b_beta * 2
+
+        if USE_G:
+            b_g = tl.load(p_g).to(tl.float32)
+            if USE_GATE_IN_KERNEL:
+                b_A = tl.load(A_log + i_hv).to(tl.float32)
+                if HAS_DT_BIAS:
+                    b_g = b_g + tl.load(dt_bias + i_hv).to(tl.float32)
+                b_g = -exp(b_A) * softplus(b_g)
+            b_h *= exp(b_g)
+
+        if USE_GK:
+            b_gk = tl.load(p_gk).to(tl.float32)
+            if STATE_V_FIRST:
+                b_h *= exp(b_gk[None, :])
+            else:
+                b_h *= exp(b_gk[:, None])
+
+        if USE_GV:
+            b_gv = tl.load(p_gv).to(tl.float32)
+            if STATE_V_FIRST:
+                b_h *= exp(b_gv[:, None])
+            else:
+                b_h *= exp(b_gv[None, :])
+
+        if STATE_V_FIRST:
+            b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1))
+            b_h += b_v[:, None] * b_k[None, :]
+            b_o = tl.sum(b_h * b_q[None, :], 1)
+        else:
+            b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0))
+            b_h += b_k[:, None] * b_v
+            b_o = tl.sum(b_h * b_q[:, None], 0)
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
+
+        p_q += H*K
+        p_k += H*K
+        p_v += HV*V
+        if USE_G:
+            p_g += HV
+        if USE_GK:
+            p_gk += HV*K
+        if USE_GV:
+            p_gv += HV*V
+        p_beta += HV * (1 if IS_BETA_HEADWISE else V)
+        p_o += HV*V
+
+    if STORE_FINAL_STATE:
+        if STATE_V_FIRST:
+            p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :]
+        else:
+            p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :]
+        tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
+
+
+def fused_recurrent_gated_delta_rule_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor | None = None,
+    gk: torch.Tensor | None = None,
+    gv: torch.Tensor | None = None,
+    beta: torch.Tensor | None = None,
+    A_log: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    use_beta_sigmoid_in_kernel: bool = False,
+    allow_neg_eigval: bool = False,
+    state_v_first: bool = False,
+    cu_seqlens: torch.LongTensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    HV = v.shape[2]
+    N = B if cu_seqlens is None else len(cu_seqlens) - 1
+    BK = triton.next_power_of_2(K)
+    BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V)
+    NV = triton.cdiv(V, BV)
+
+    o = torch.empty_like(v)
+    if output_final_state:
+        if state_v_first:
+            final_state = q.new_empty(N, HV, V, K, dtype=torch.float32)
+        else:
+            final_state = q.new_empty(N, HV, K, V, dtype=torch.float32)
+    else:
+        final_state = None
+
+    grid = (NV * N * HV,)
+    fused_recurrent_gated_delta_rule_fwd_kernel[grid](
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        gk=gk,
+        gv=gv,
+        beta=beta,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        o=o,
+        h0=initial_state,
+        ht=final_state,
+        cu_seqlens=cu_seqlens,
+        scale=scale,
+        T=T,
+        H=H,
+        HV=HV,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        IS_BETA_HEADWISE=beta.ndim != v.ndim,
+        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,
+        ALLOW_NEG_EIGVAL=allow_neg_eigval,
+        STATE_V_FIRST=state_v_first,
+        num_warps=1,
+        num_stages=3,
+    )
+    return o, final_state
+
+
+class FusedRecurrentFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    def forward(
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor | None = None,
+        gk: torch.Tensor | None = None,
+        gv: torch.Tensor | None = None,
+        beta: torch.Tensor | None = None,
+        A_log: torch.Tensor | None = None,
+        dt_bias: torch.Tensor | None = None,
+        scale: float = None,
+        initial_state: torch.Tensor = None,
+        output_final_state: bool = False,
+        use_qk_l2norm_in_kernel: bool = False,
+        use_beta_sigmoid_in_kernel: bool = False,
+        allow_neg_eigval: bool = False,
+        state_v_first: bool = False,
+        cu_seqlens: torch.LongTensor | None = None,
+    ):
+        o, final_state = fused_recurrent_gated_delta_rule_fwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            gk=gk,
+            gv=gv,
+            beta=beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            scale=scale,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+            use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
+            allow_neg_eigval=allow_neg_eigval,
+            state_v_first=state_v_first,
+            cu_seqlens=cu_seqlens,
+        )
+
+        return o, final_state
+
+    @staticmethod
+    @input_guard
+    def backward(ctx, do, dht):
+        raise NotImplementedError(
+            "Backward pass is not implemented yet and we do not have plans to implement it "
+            "because we haven't figured out how to compute dg without materializing the full "
+            "hidden states for all time steps.",
+        )
+
+
+def fused_recurrent_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor | None = None,
+    gk: torch.Tensor | None = None,
+    gv: torch.Tensor | None = None,
+    beta: torch.Tensor | None = None,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    use_gate_in_kernel: bool = False,
+    A_log: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    use_beta_sigmoid_in_kernel: bool = False,
+    allow_neg_eigval: bool = False,
+    state_v_first: bool = False,
+    cu_seqlens: torch.LongTensor | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    r"""
+    Args:
+        q (torch.Tensor):
+            queries of shape `[B, T, H, K]`.
+        k (torch.Tensor):
+            keys of shape `[B, T, H, K]`.
+        v (torch.Tensor):
+            values of shape `[B, T, HV, V]`.
+            GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.
+        g (torch.Tensor):
+            g (decays) of shape `[B, T, HV]`. Default: `None`.
+            When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay).
+            When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses
+            `-exp(A_log) * softplus(g + dt_bias)` internally per step.
+        gk (torch.Tensor):
+            gk (decays) of shape `[B, T, HV, K]`. Default: `None`.
+        gv (torch.Tensor):
+            gv (decays) of shape `[B, T, HV, V]`. Default: `None`.
+        beta (torch.Tensor):
+            betas of shape `[B, T, HV]`.
+        scale (Optional[float]):
+            Scale factor for the RetNet attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `[N, HV, K, V]` for `N` input sequences.
+            For equal-length input sequences, `N` equals the batch size `B`.
+            Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
+        use_qk_l2norm_in_kernel (Optional[bool]):
+            Whether to use L2 normalization in the kernel. Default: `False`.
+        use_gate_in_kernel (bool):
+            Whether to compute the log-space GDN decay internally.
+            When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses
+            gate activation into the recurrence. Default: `False`.
+        A_log (Optional[torch.Tensor]):
+            Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.
+        dt_bias (Optional[torch.Tensor]):
+            Bias added to `g` before activation, of shape `[HV]`.
+            Only used when `use_gate_in_kernel=True`.
+        use_beta_sigmoid_in_kernel (Optional[bool]):
+            Whether to apply `torch.sigmoid(beta)` inside the kernel.
+            - If `True`, the passed `beta` acts as the raw beta logits.
+            - If `False`, `beta` is expected to already be in post-sigmoid space.
+            Default: `False`.
+        allow_neg_eigval (Optional[bool]):
+            Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.
+            Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case
+            the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.
+        state_v_first (Optional[bool]):
+            Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.
+        cu_seqlens (torch.LongTensor):
+            Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
+            consistent with the FlashAttention API.
+
+    Returns:
+        o (torch.Tensor):
+            Outputs of shape `[B, T, HV, V]`.
+        final_state (torch.Tensor):
+            Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
+
+    Examples::
+        >>> import torch
+        >>> import torch.nn.functional as F
+        >>> from einops import rearrange
+        >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
+        # inputs with equal lengths
+        >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
+        >>> q = torch.randn(B, T, H, K, device='cuda')
+        >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)
+        >>> v = torch.randn(B, T, HV, V, device='cuda')
+        >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
+        >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
+        >>> h0 = torch.randn(B, HV, K, V, device='cuda')
+        >>> o, ht = fused_gated_recurrent_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True
+        )
+        # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
+        >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))
+        # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
+        >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
+        >>> o, ht = fused_gated_recurrent_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True,
+            cu_seqlens=cu_seqlens
+        )
+    """
+    if 'transpose_state_layout' in kwargs:
+        if state_v_first:
+            raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.")
+        warnings.warn(
+            "`transpose_state_layout` is deprecated and renamed to `state_v_first`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        state_v_first = kwargs.pop('transpose_state_layout')
+
+    if cu_seqlens is not None:
+        if q.shape[0] != 1:
+            raise ValueError(
+                f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
+                f"Please flatten variable-length inputs before processing.",
+            )
+        if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
+            raise ValueError(
+                f"The number of initial states is expected to be equal to the number of input sequences, "
+                f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.",
+            )
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    if beta is None:
+        beta = torch.ones_like(q[..., 0])
+    if use_gate_in_kernel:
+        if A_log is None:
+            raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True`.")
+        if g is None:
+            raise ValueError("`g` (raw pre-activation) must be provided when `use_gate_in_kernel=True`.")
+    else:
+        A_log = None
+        dt_bias = None
+    if allow_neg_eigval and not use_beta_sigmoid_in_kernel:
+        raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.")
+
+    o, final_state = FusedRecurrentFunction.apply(
+        q,
+        k,
+        v,
+        g,
+        gk,
+        gv,
+        beta,
+        A_log,
+        dt_bias,
+        scale,
+        initial_state,
+        output_final_state,
+        use_qk_l2norm_in_kernel,
+        use_beta_sigmoid_in_kernel,
+        allow_neg_eigval,
+        state_v_first,
+        cu_seqlens,
+    )
+    return o, final_state
+
+
+fused_recurrent_gdn = fused_recurrent_gated_delta_rule

ex_engine/fla_kernels/gated_delta_rule/gate.py

@@ -0,0 +1,344 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from fla.ops.backends import dispatch
+from fla.ops.utils.cache import fla_cache_autotune
+from fla.ops.utils.index import prepare_chunk_indices
+from fla.ops.utils.op import exp
+from fla.ops.utils.softplus import softplus
+from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard
+
+
+def naive_gdn_gate(
+    g: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor | None = None,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """
+    Torch reference implementation for GDN gate computation.
+
+    Computes: ``g = -A_log.exp() * softplus(g + dt_bias)``
+
+    Args:
+        g (torch.Tensor):
+            Input tensor of shape `[..., HV]`.
+        A_log (torch.Tensor):
+            Decay parameter tensor with `HV` elements.
+        dt_bias (torch.Tensor | None):
+            Optional bias tensor added to `g` before activation, shape `[HV]`.
+
+    Returns:
+        Output tensor of shape `[..., HV]`.
+    """
+    g = g.float()
+    if dt_bias is not None:
+        g = g + dt_bias.float()
+    return (-A_log.float().exp() * F.softplus(g)).to(output_dtype)
+
+
+@triton.heuristics({
+    'HAS_BIAS': lambda args: args['dt_bias'] is not None,
+    'HAS_SCALE': lambda args: args['scale'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@fla_cache_autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['H', 'BT', 'IS_VARLEN', 'REVERSE'],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=['T'])
+def gdn_gate_chunk_cumsum_scalar_kernel(
+    g,
+    A_log,
+    dt_bias,
+    o,
+    scale,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    REVERSE: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    HAS_SCALE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    i_b, i_h = i_bh // H, i_bh % H
+
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    o_t = i_t * BT + tl.arange(0, BT)
+    m_t = o_t < T
+    p_g = g + bos * H + i_h + o_t * H
+    p_o = o + bos * H + i_h + o_t * H
+
+    b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
+    if HAS_BIAS:
+        b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
+    b_A = tl.load(A_log + i_h).to(tl.float32)
+    b_gate = -exp(b_A) * softplus(b_g)
+
+    b_o = tl.cumsum(b_gate, axis=0)
+    if REVERSE:
+        b_z = tl.sum(b_gate, axis=0)
+        b_o = -b_o + b_z[None] + b_gate
+    if HAS_SCALE:
+        b_o *= scale
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t)
+
+
+@triton.heuristics({
+    'HAS_BIAS': lambda args: args['dt_bias'] is not None,
+})
+@fla_cache_autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['H', 'BT'],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=['T'])
+def gdn_gate_bwd_kernel(
+    g,
+    A_log,
+    dt_bias,
+    dyg,
+    dg,
+    dA,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+):
+    i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1)
+
+    b_A = tl.load(A_log + i_h).to(tl.float32)
+
+    o_t = i_t * BT + tl.arange(0, BT)
+    m_t = o_t < T
+    p_g = g + i_h + o_t * H
+    p_dg = dg + i_h + o_t * H
+    p_dyg = dyg + i_h + o_t * H
+
+    b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
+    b_dyg = tl.load(p_dyg, mask=m_t, other=0.0).to(tl.float32)
+
+    if HAS_BIAS:
+        b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
+
+    # gate = -exp(A_log) * softplus(g + bias)
+    # d(gate)/d(g) = -exp(A_log) * sigmoid(g + bias)   (softplus' = sigmoid)
+    # d(gate)/d(A_log) = -exp(A_log) * softplus(g + bias) = gate
+    b_neg_expA = -exp(b_A)
+    b_yg = b_neg_expA * softplus(b_g)
+    b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g))
+    b_dA = tl.sum(b_dyg * b_yg, 0)
+
+    tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t)
+    tl.store(dA + i_t * H + i_h, b_dA)
+
+
+@input_guard
+@dispatch('gated_delta_rule')
+def gdn_gate_chunk_cumsum(
+    g: torch.Tensor,
+    A_log: torch.Tensor,
+    chunk_size: int,
+    scale: float = None,
+    dt_bias: torch.Tensor | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+    output_dtype: torch.dtype | None = torch.float,
+) -> torch.Tensor:
+    B, T, H = g.shape
+    BT = chunk_size
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    o = torch.empty_like(g, dtype=output_dtype or g.dtype)
+    gdn_gate_chunk_cumsum_scalar_kernel[(NT, B * H)](
+        g=g,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        o=o,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        BT=BT,
+        REVERSE=False,
+    )
+    return o
+
+
+@dispatch('gated_delta_rule')
+def gdn_gate_bwd(
+    g: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor | None,
+    dyg: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    H = g.shape[-1]
+    T = g.numel() // H
+    BT = 32
+    NT = triton.cdiv(T, BT)
+
+    dg = torch.empty_like(g, dtype=torch.float32)
+    dA = A_log.new_empty(NT, H, dtype=torch.float32)
+
+    gdn_gate_bwd_kernel[(NT, H)](
+        g=g,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        dyg=dyg,
+        dg=dg,
+        dA=dA,
+        T=T,
+        H=H,
+        BT=BT,
+    )
+
+    dg = dg.view_as(g).type_as(g)
+    dA = dA.sum(0).view_as(A_log).type_as(A_log)
+    dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None
+
+    return dg, dA, dbias
+
+
+@triton.heuristics({
+    'HAS_BIAS': lambda args: args['dt_bias'] is not None,
+})
+@fla_cache_autotune(
+    configs=[
+        triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages)
+        for BT in [32, 64, 128]
+        for num_warps in [1, 2, 4, 8]
+        for num_stages in [2, 3]
+    ],
+    key=['H'],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=['T'])
+def gdn_gate_fwd_kernel(
+    g,
+    A_log,
+    dt_bias,
+    yg,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+):
+    i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1)
+
+    b_A = tl.load(A_log + i_h).to(tl.float32)
+
+    o_t = i_t * BT + tl.arange(0, BT)
+    m_t = o_t < T
+    p_g = g + i_h + o_t * H
+    p_yg = yg + i_h + o_t * H
+    b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
+    if HAS_BIAS:
+        b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
+    b_yg = -exp(b_A) * softplus(b_g)
+    tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_t)
+
+
+@dispatch('gated_delta_rule')
+def gdn_gate_fwd(
+    g: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor | None = None,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    H = g.shape[-1]
+    T = g.numel() // H
+
+    yg = torch.empty_like(g, dtype=output_dtype)
+
+    def grid(meta):
+        return (triton.cdiv(T, meta['BT']), H)
+
+    gdn_gate_fwd_kernel[grid](
+        g=g,
+        A_log=A_log,
+        dt_bias=dt_bias,
+        yg=yg,
+        T=T,
+        H=H,
+    )
+    return yg
+
+
+class GDNGateFunction(torch.autograd.Function):
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(
+        ctx,
+        g: torch.Tensor,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor | None = None,
+        output_dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor:
+        yg = gdn_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, output_dtype=output_dtype)
+        ctx.save_for_backward(g, A_log, dt_bias)
+        return yg
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(ctx, dyg: torch.Tensor):
+        g, A_log, dt_bias = ctx.saved_tensors
+        dg, dA, dbias = gdn_gate_bwd(g=g, A_log=A_log, dt_bias=dt_bias, dyg=dyg)
+        return dg, dA, dbias, None
+
+
+@torch.compiler.disable
+def fused_gdn_gate(
+    g: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor | None = None,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    r"""
+    Fused GDN gate computation with autograd support.
+
+    Computes: ``g = -A_log.exp() * softplus(g + dt_bias)``
+
+    Args:
+        g (torch.Tensor):
+            Input tensor of shape `[..., HV]`.
+        A_log (torch.Tensor):
+            Decay parameter tensor with `HV` elements.
+        dt_bias (torch.Tensor | None):
+            Optional bias tensor added to `g` before activation, shape `[HV]`.
+        output_dtype (torch.dtype):
+            The dtype of the output tensor. Default: `torch.float32`.
+
+    Returns:
+        Output tensor of shape `[..., HV]`.
+    """
+    return GDNGateFunction.apply(g, A_log, dt_bias, output_dtype)

ex_engine/fla_kernels/gated_delta_rule/naive.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+
+def naive_recurrent_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    g: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    """
+    Reference PyTorch implementation of recurrent gated delta rule.
+
+    Args:
+        q: [B, T, H, K]
+        k: [B, T, H, K]
+        v: [B, T, H, V]
+        beta: [B, T, H]
+        g: [B, T, H]
+        scale: float, optional
+        initial_state: [B, H, K, V], optional
+        output_final_state: bool
+
+    Returns:
+        o: [B, T, H, V]
+        final_state: [B, H, K, V] if output_final_state else None
+    """
+    q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    o = torch.zeros(B, H, T, V).to(v)
+    h = torch.zeros(B, H, K, V).to(v)
+    if initial_state is not None:
+        h = initial_state.to(torch.float32)
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+    q = q * scale
+
+    for i in range(T):
+        b_q = q[:, :, i]
+        b_k = k[:, :, i]
+        b_v = v[:, :, i].clone()
+        h = h.clone() * g[:, :, i].exp()[..., None, None]
+        b_beta = beta[:, :, i]
+        b_v = b_v - (h.clone() * b_k[..., None]).sum(-2)
+        b_v = b_v * b_beta[..., None]
+        h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
+        o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h)
+
+    if not output_final_state:
+        h = None
+    o = o.transpose(1, 2).contiguous()
+    return o, h
+
+
+def naive_chunk_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = 64,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    """
+    Reference PyTorch implementation of chunk gated delta rule.
+
+    Args:
+        q: [B, T, H, K]
+        k: [B, T, H, K]
+        v: [B, T, H, V]
+        g: [B, T, H]
+        beta: [B, T, H]
+        chunk_size: int
+        scale: float, optional
+        initial_state: [B, H, K, V], optional
+        output_final_state: bool
+
+    Returns:
+        o: [B, T, H, V]
+        final_state: [B, H, K, V] if output_final_state else None
+    """
+    BT = chunk_size
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+
+    q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
+
+    T = q.shape[-2]
+    pad_len = (BT - (T % BT)) % BT
+    if pad_len > 0:
+        q = F.pad(q, (0, 0, 0, pad_len))
+        k = F.pad(k, (0, 0, 0, pad_len))
+        v = F.pad(v, (0, 0, 0, pad_len))
+        beta = F.pad(beta, (0, pad_len))
+        g = F.pad(g, (0, pad_len))
+
+    q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g])
+    decay = g
+    chunk_size = BT
+    b, h, l, d_k = q.shape
+    d_v = v.shape[-1]
+    q = q * scale
+    v = v * beta[..., None]
+    k_beta = k * beta[..., None]
+    assert l % chunk_size == 0
+
+    # note that diagonal is masked.
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
+    q, k, v, k_beta, decay = map(
+        lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size),
+        [q, k, v, k_beta, decay.unsqueeze(-1)],
+    )
+    decay = decay.squeeze(-1).cumsum(-1)
+    decay_exp = decay.exp()[..., None]
+    L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0)
+    for i in range(1, chunk_size):
+        attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
+    attn = attn
+    k_cumsum = attn @ v
+    k_cumdecay = attn @ (k_beta * decay_exp)
+    v = k_cumsum
+
+    S = k.new_zeros(b, h, d_k, d_v)
+    if initial_state is not None:
+        S = initial_state.to(torch.float32)
+
+    o = torch.zeros_like(v)
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
+    for i in range(0, l // chunk_size):
+        q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
+        attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
+        v_prime = (k_cumdecay[:, :, i]) @ S
+        v_new = v_i - v_prime
+        o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
+        o[:, :, i] = o_inter + attn @ v_new
+        S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()
+                                                        [..., None]).transpose(-1, -2) @ v_new
+    if not output_final_state:
+        S = None
+
+    # unpad
+    o = rearrange(o, 'b h n c d -> b h (n c) d')
+    o = o[:, :, :T]
+    o = o.transpose(1, 2)
+    return o, S

ex_engine/fla_kernels/gated_delta_rule/wy_fast.py

@@ -0,0 +1,351 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import torch
+import triton
+import triton.language as tl
+
+from fla.ops.backends import dispatch
+from fla.ops.utils import prepare_chunk_indices
+from fla.ops.utils.cache import fla_cache_autotune
+from fla.ops.utils.op import exp2
+from fla.utils import IS_INTEL, IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem
+
+# Blackwell can select unstable Triton configs for prepare_wy_repr_bwd_kernel
+# during autotuning (see #913). Restrict it to the config that has been
+# validated on B200 until the wider config space is re-validated.
+PREPARE_WY_REPR_BWD_NUM_WARPS = [2] if IS_NVIDIA_BLACKWELL else [2, 4]
+PREPARE_WY_REPR_BWD_NUM_STAGES = [4] if IS_NVIDIA_BLACKWELL else [2, 3, 4]
+
+# Intel keeps scaling past the warp counts NVIDIA prefers: 16 warps is ~1.3x faster
+# than 8 for recompute_w_u.
+RECOMPUTE_W_U_NUM_WARPS = [2, 4, 8, 16] if IS_INTEL else [2, 4, 8]
+
+
+@triton.heuristics({
+    'USE_G': lambda args: args['g'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@fla_cache_autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in RECOMPUTE_W_U_NUM_WARPS
+        for num_stages in [2, 3, 4]
+    ],
+    key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=['T'])
+def recompute_w_u_fwd_kernel(
+    k,
+    v,
+    beta,
+    w,
+    u,
+    A,
+    g,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    i_b, i_h = i_bh // HV, i_bh % HV
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+    o_t = i_t * BT + tl.arange(0, BT)
+    o_A = tl.arange(0, BT)
+    m_t = o_t < T
+    m_A = m_t[:, None] & (o_A[None, :] < BT)
+    p_b = beta + bos*HV + i_h + o_t * HV
+    b_b = tl.load(p_b, mask=m_t, other=0.0)
+
+    p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (HV*BT) + o_A[None, :]
+    b_A = tl.load(p_A, mask=m_A, other=0.0)
+
+    for i_v in range(tl.cdiv(V, BV)):
+        o_v = i_v * BV + tl.arange(0, BV)
+        m_v = m_t[:, None] & (o_v[None, :] < V)
+        p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
+        p_u = u + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
+        b_v = tl.load(p_v, mask=m_v, other=0.0)
+        b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
+        b_u = tl.dot(b_A, b_vb, allow_tf32=False)
+        tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v)
+
+    if USE_G:
+        p_g = g + (bos*HV + i_h) + o_t * HV
+        b_g = exp2(tl.load(p_g, mask=m_t, other=0.0))
+
+    for i_k in range(tl.cdiv(K, BK)):
+        o_k = i_k * BK + tl.arange(0, BK)
+        m_k = m_t[:, None] & (o_k[None, :] < K)
+        p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
+        p_w = w + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
+        b_k = tl.load(p_k, mask=m_k, other=0.0)
+        b_kb = b_k * b_b[:, None]
+        if USE_G:
+            b_kb *= b_g[:, None]
+        b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
+        tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k)
+
+
+@triton.heuristics({
+    'USE_G': lambda args: args['g'] is not None,
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
+})
+@fla_cache_autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps, num_stages=num_stages)
+        for num_warps in PREPARE_WY_REPR_BWD_NUM_WARPS
+        for num_stages in PREPARE_WY_REPR_BWD_NUM_STAGES
+    ],
+    key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=['T'])
+def prepare_wy_repr_bwd_kernel(
+    k,
+    v,
+    beta,
+    g,
+    A,
+    dw,
+    du,
+    dk,
+    dv,
+    db,
+    dg,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    i_b, i_h = i_bh // HV, i_bh % HV
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    o_t = i_t * BT + tl.arange(0, BT)
+    o_A = tl.arange(0, BT)
+    m_t = o_t < T
+    m_AT = (o_A[:, None] < BT) & m_t[None, :]
+    p_b = beta + (bos*HV + i_h) + o_t * HV
+    p_db = db + (bos*HV + i_h) + o_t * HV
+    p_A = A + (bos*HV + i_h) * BT + o_A[:, None] + o_t[None, :] * (HV*BT)
+
+    b_b = tl.load(p_b, mask=m_t, other=0.0)
+    b_db = tl.zeros([BT], dtype=tl.float32)
+    b_A = tl.load(p_A, mask=m_AT, other=0.0)
+    b_dA = tl.zeros([BT, BT], dtype=tl.float32)
+
+    if USE_G:
+        p_g = g + (bos*HV + i_h) + o_t * HV
+        b_g = tl.load(p_g, mask=m_t, other=0.0)
+        b_g_exp = exp2(b_g)
+        b_dg = tl.zeros([BT], dtype=tl.float32)
+
+    for i_k in range(tl.cdiv(K, BK)):
+        o_k = i_k * BK + tl.arange(0, BK)
+        m_k = m_t[:, None] & (o_k[None, :] < K)
+        p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
+        p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
+        p_dw = dw + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
+        # [BT, BK]
+        b_k = tl.load(p_k, mask=m_k, other=0.0)
+        if USE_G:
+            b_kbg = b_k * (b_b * b_g_exp)[:, None]
+        else:
+            b_kbg = b_k * b_b[:, None]
+        b_dw = tl.load(p_dw, mask=m_k, other=0.0)
+
+        b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype))
+        b_dkbg = tl.dot(b_A, b_dw)
+        if USE_G:
+            b_dk = b_dkbg * (b_g_exp * b_b)[:, None]
+            b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1)
+            b_dg += tl.sum(b_dkbg * b_kbg, 1)
+        else:
+            b_dk = b_dkbg * b_b[:, None]
+            b_db += tl.sum(b_dkbg * b_k, 1)
+        tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k)
+
+    for i_v in range(tl.cdiv(V, BV)):
+        o_v = i_v * BV + tl.arange(0, BV)
+        m_v = m_t[:, None] & (o_v[None, :] < V)
+        p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
+        p_dv = dv + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
+        p_du = du + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
+        b_v = tl.load(p_v, mask=m_v, other=0.0)
+        b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
+        b_du = tl.load(p_du, mask=m_v, other=0.0)
+        b_dA += tl.dot(b_du, tl.trans(b_vb))
+        b_dvb = tl.dot(b_A, b_du)
+        b_dv = b_dvb * b_b[:, None]
+        b_db += tl.sum(b_dvb * b_v, 1)
+        tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v)
+
+    m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
+    b_dA = tl.where(m_A, b_dA, 0)
+    b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
+    b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
+
+    if USE_G:
+        b_dA *= exp2(b_g[:, None] - b_g[None, :])
+
+    b_A = tl.zeros([BT, BT], dtype=tl.float32)
+    b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty)
+
+    tl.debug_barrier()
+    for i_k in range(tl.cdiv(K, BK)):
+        o_k = i_k * BK + tl.arange(0, BK)
+        m_k = m_t[:, None] & (o_k[None, :] < K)
+        p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
+        p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
+        b_k = tl.load(p_k, mask=m_k, other=0.0)
+        b_kt = tl.trans(b_k)
+        b_kb = b_k * b_b[:, None]
+
+        b_A += tl.dot(b_k, b_kt)
+        b_dkb = tl.dot(b_dA, b_k)
+        b_db += tl.sum(b_dkb * b_k, 1)
+        b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA))
+        b_dk += tl.load(p_dk, mask=m_k, other=0.0)
+
+        tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k)
+    tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t)
+
+    b_A *= b_b[:, None]
+    if USE_G:
+        b_AdA = b_dA * b_A
+        p_dg = dg + (bos*HV + i_h) + o_t * HV
+        b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0)
+        tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t)
+
+
+@dispatch('gated_delta_rule')
+def recompute_w_u_fwd(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    g: torch.Tensor | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2]
+    BT = A.shape[-1]
+    BK = 64
+    BV = 64
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    w = k.new_empty(B, T, HV, K)
+    u = torch.empty_like(v)
+    recompute_w_u_fwd_kernel[(NT, B*HV)](
+        k=k,
+        v=v,
+        beta=beta,
+        w=w,
+        u=u,
+        A=A,
+        g=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        HV=HV,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+    return w, u
+
+
+@dispatch('gated_delta_rule')
+def prepare_wy_repr_bwd(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    dw: torch.Tensor,
+    du: torch.Tensor,
+    g: torch.Tensor = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_indices: torch.LongTensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2]
+    BT = A.shape[-1]
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    CONST_TILING = 64 if check_shared_mem() else 32
+    BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING)
+    BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING)
+
+    dk = k.new_empty(B, T, HV, K)
+    dv = torch.empty_like(v)
+    dg = torch.empty_like(g) if g is not None else None
+    db = torch.empty_like(beta)
+    prepare_wy_repr_bwd_kernel[(NT, B * HV)](
+        k=k,
+        v=v,
+        beta=beta,
+        g=g,
+        A=A,
+        dw=dw,
+        du=du,
+        dk=dk,
+        dv=dv,
+        db=db,
+        dg=dg,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        HV=HV,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+    if H != HV:
+        dk = dk.view(B, T, H, HV // H, K).sum(3)
+    return dk, dv, db, dg
+
+
+fwd_recompute_w_u = recompute_w_u_fwd
+bwd_prepare_wy_repr = prepare_wy_repr_bwd

ex_engine/fla_kernels/utils/__init__.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+from .csr import prepare_block_csr
+from .cumsum import (
+    chunk_global_cumsum,
+    chunk_global_cumsum_scalar,
+    chunk_global_cumsum_vector,
+    chunk_local_cumsum,
+    chunk_local_cumsum_scalar,
+    chunk_local_cumsum_vector,
+)
+from .index import (
+    get_max_num_splits,
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+    prepare_cu_seqlens_from_lens,
+    prepare_cu_seqlens_from_mask,
+    prepare_lens,
+    prepare_lens_from_mask,
+    prepare_position_ids,
+    prepare_sequence_ids,
+    prepare_token_indices,
+)
+from .logsumexp import logsumexp_fwd
+from .matmul import addmm, matmul
+from .pack import pack_sequence, unpack_sequence
+from .pooling import mean_pooling
+from .softmax import softmax_bwd, softmax_fwd
+from .softplus import softplus
+from .solve_tril import solve_tril
+
+__all__ = [
+    "addmm",
+    "chunk_global_cumsum",
+    "chunk_global_cumsum_scalar",
+    "chunk_global_cumsum_vector",
+    "chunk_local_cumsum",
+    "chunk_local_cumsum_scalar",
+    "chunk_local_cumsum_vector",
+    "get_max_num_splits",
+    "logsumexp_fwd",
+    "matmul",
+    "mean_pooling",
+    "pack_sequence",
+    "prepare_block_csr",
+    "prepare_chunk_indices",
+    "prepare_chunk_offsets",
+    "prepare_cu_seqlens_from_lens",
+    "prepare_cu_seqlens_from_mask",
+    "prepare_lens",
+    "prepare_lens_from_mask",
+    "prepare_position_ids",
+    "prepare_sequence_ids",
+    "prepare_token_indices",
+    "softmax_bwd",
+    "softmax_fwd",
+    "softplus",
+    "solve_tril",
+    "unpack_sequence",
+]

ex_engine/fla_kernels/utils/cache.py

@@ -0,0 +1,449 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import dataclasses
+import enum
+import json
+import logging
+import os
+import re
+from functools import cache, lru_cache
+from pathlib import Path
+from typing import Any
+
+import torch
+import triton
+from packaging import version
+from triton.runtime.autotuner import Autotuner
+
+TRITON_ABOVE_3_5_1 = version.parse(triton.__version__) >= version.parse("3.5.1")
+TRITON_ABOVE_3_4_0 = version.parse(triton.__version__) >= version.parse("3.4.0")
+
+
+class FlaCacheMode(enum.Enum):
+    """Controls how FLA loads kernel configs from its config cache (FLA_CACHE_MODE env var).
+
+    DISABLED    — skip all cache lookups, always fall back to Triton autotune (default when FLA_CACHE_MODE is unset)
+    STRICT      — exact key match only; falls back to Triton autotune if no match
+    FUZZY       — exact key match → fuzzy key match; falls back to Triton autotune if no match
+    FULL        — exact key match → fuzzy key match → default_config fallback
+    DEFAULT     — use only the top-level default_config field, skip key-based lookup
+    ALWAYS      — like DEFAULT, but re-reads config files on every kernel call;
+                  useful for debugging: edit default_config in a JSON file and the next
+                  kernel call picks it up without restarting the process
+    """
+    DISABLED = "disabled"
+    STRICT = "strict"
+    FUZZY = "fuzzy"
+    FULL = "full"
+    DEFAULT = "default"
+    ALWAYS = "always"
+
+    def uses_default_config(self) -> bool:
+        """Return True for modes that may fall back to default_config (FULL, DEFAULT, ALWAYS)."""
+        return self in (FlaCacheMode.FULL, FlaCacheMode.DEFAULT, FlaCacheMode.ALWAYS)
+
+    @classmethod
+    def from_env(cls) -> "FlaCacheMode":
+        mode_str = os.environ.get("FLA_CACHE_MODE", cls.DISABLED.value)
+        try:
+            return cls(mode_str)
+        except ValueError:
+            valid = [m.value for m in cls]
+            raise ValueError(
+                f"Invalid FLA_CACHE_MODE={mode_str!r}. Valid values: {valid}"
+            ) from None
+
+
+FLA_CACHE_MODE: FlaCacheMode = FlaCacheMode.from_env()
+logger = logging.getLogger(__name__)
+
+
+def sanitize_gpu_name(gpu_name: str) -> str:
+    sanitized = re.sub(r"[^0-9A-Za-z]+", "_", gpu_name)
+    sanitized = sanitized.strip("_")
+    return sanitized or "unknown_gpu"
+
+
+@lru_cache(maxsize=1)
+def get_gpu_info():
+    """Get GPU model information.
+
+    This function detects the GPU model and returns a sanitized string identifier.
+    It prioritizes FLA_GPU_NAME environment variable if set, then detects from
+    available hardware (CUDA, ROCm, Intel GPU, or CPU).
+    """
+    # Check if GPU name is overridden via environment variable
+    gpu_name = None
+    # Check if GPU name is overridden via environment variable
+    if "FLA_GPU_NAME" in os.environ:
+        gpu_name = os.environ["FLA_GPU_NAME"]
+    # Try to get device name based on availability
+    elif torch.cuda.is_available():
+        # Works for both NVIDIA and AMD GPUs (ROCm)
+        gpu_name = torch.cuda.get_device_name(0)
+    elif hasattr(torch, 'xpu') and torch.xpu.is_available():
+        gpu_name = torch.xpu.get_device_name(0)
+
+    if gpu_name:
+        return sanitize_gpu_name(gpu_name)
+
+    # Default to CPU if no GPU available
+    return "cpu"
+
+
+def get_fla_config_dir() -> Path:
+    """Get FLA's configs directory.
+
+    The directory can be overridden by setting the FLA_CONFIG_DIR environment variable.
+    If set, configs will be loaded directly from $FLA_CONFIG_DIR/. Otherwise FLA
+    falls back to the default fla/configs/{GPU}/ directory in the project.
+    """
+    # Check if custom config dir is set via environment variable
+    if "FLA_CONFIG_DIR" in os.environ:
+        return Path(os.environ["FLA_CONFIG_DIR"])
+
+    # Default: project_dir/fla/configs/{GPU}/
+    project_dir = Path(__file__).parent.parent.parent
+    return project_dir / "configs" / get_gpu_info()
+
+
+@dataclasses.dataclass(frozen=True)
+class AutotuneKey:
+    """Autotune key with exact/fuzzy matching, serialization, and construction helpers."""
+    autotune_key: tuple[Any, ...]
+
+    @staticmethod
+    def normalize_autotune_key(value: Any) -> Any:
+        if isinstance(value, (list, tuple)):
+            return [AutotuneKey.normalize_autotune_key(v) for v in value]
+        if isinstance(value, dict):
+            return {k: AutotuneKey.normalize_autotune_key(v) for k, v in value.items()}
+        return value
+
+    @staticmethod
+    def serialize(key: Any) -> str:
+        return json.dumps(AutotuneKey.normalize_autotune_key(key), separators=(",", ":"), sort_keys=True)
+
+    @staticmethod
+    def key_hash(key: Any) -> str:
+        import hashlib
+        return hashlib.md5(AutotuneKey.serialize(key).encode()).hexdigest()
+
+    @staticmethod
+    def is_numeric(value: Any) -> bool:
+        return isinstance(value, (int, float)) and not isinstance(value, bool)
+
+    @staticmethod
+    def keys_fuzzy_match(cached_key: Any, requested_key: Any) -> bool:
+        # Fuzzy match: numeric leaves are compatible regardless of their actual numeric values
+        # (e.g. a config tuned for seq_len=1024 can apply to seq_len=2048).
+        # Structure (type, length, dict keys) must still match exactly.
+        if AutotuneKey.is_numeric(cached_key) and AutotuneKey.is_numeric(requested_key):
+            return True
+        if isinstance(cached_key, (list, tuple)) and isinstance(requested_key, (list, tuple)):
+            return len(cached_key) == len(requested_key) and all(
+                AutotuneKey.keys_fuzzy_match(c, r) for c, r in zip(cached_key, requested_key)
+            )
+        if isinstance(cached_key, dict) and isinstance(requested_key, dict):
+            return cached_key.keys() == requested_key.keys() and all(
+                AutotuneKey.keys_fuzzy_match(cached_key[k], requested_key[k]) for k in cached_key
+            )
+        return cached_key == requested_key
+
+    @classmethod
+    def build(
+        cls,
+        arg_names: list[str],
+        key_names: list[str],
+        positional_args: tuple[Any, ...],
+        runtime_kwargs: dict[str, Any],
+    ) -> "AutotuneKey":
+        named_args = dict(zip(arg_names, positional_args))
+        all_args = {**named_args, **runtime_kwargs}
+        tracked_args = {k: v for (k, v) in all_args.items() if k in arg_names}
+        tuning_key = [tracked_args[name] for name in key_names if name in tracked_args]
+        for arg in tracked_args.values():
+            if hasattr(arg, "dtype"):
+                tuning_key.append(str(arg.dtype))
+        return cls(autotune_key=tuple(tuning_key))
+
+    def exact_matches(self, entry_key: Any) -> bool:
+        return self.serialize(self.autotune_key) == self.serialize(entry_key)
+
+    def fuzzy_matches(self, entry_key: Any) -> bool:
+        self_normalized = self.normalize_autotune_key(self.autotune_key)
+        entry_normalized = self.normalize_autotune_key(entry_key)
+        return (
+            isinstance(self_normalized, list)
+            and isinstance(entry_normalized, list)
+            and len(self_normalized) == len(entry_normalized)
+            and AutotuneKey.keys_fuzzy_match(self_normalized, entry_normalized)
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class KernelConfigFile:
+    """Validated in-memory representation of a {kernel_name}.json config file."""
+    kernel_name: str | None
+    triton_version: str | None
+    autotune_entries: dict[str, dict[str, Any]] | None
+    default_config: dict[str, Any] | None
+
+    @classmethod
+    def from_dict(cls, config_file: Path, data: Any) -> "KernelConfigFile | None":
+        """Parse and validate a raw JSON dict. Returns None (with a warning) if malformed."""
+        def fail(msg, *args):
+            logger.warning(msg, *args)
+            raise ValueError
+
+        try:
+            if not isinstance(data, dict):
+                fail("Malformed config %s: root is %s, expected dict", config_file, type(data).__name__)
+            raw_entries = data.get("autotune_entries")
+            entries: dict[str, dict[str, Any]] | None = None
+            if raw_entries is not None:
+                if not isinstance(raw_entries, dict):
+                    fail("Malformed config %s: 'autotune_entries' is %s, expected dict",
+                         config_file, type(raw_entries).__name__)
+                for h, entry in raw_entries.items():
+                    if not isinstance(entry, dict):
+                        fail("Malformed config %s: autotune_entries[%r] is %s, expected dict",
+                             config_file, h, type(entry).__name__)
+                    if not isinstance(entry.get("config"), dict):
+                        fail("Malformed config %s: autotune_entries[%r] missing valid 'config' field", config_file, h)
+                entries = raw_entries
+            default_config = data.get("default_config")
+            if default_config is not None and not isinstance(default_config, dict):
+                fail("Malformed config %s: 'default_config' is %s, expected dict", config_file, type(default_config).__name__)
+            return cls(
+                kernel_name=data.get("kernel_name"),
+                triton_version=data.get("triton_version"),
+                autotune_entries=entries,
+                default_config=default_config,
+            )
+        except ValueError:
+            return None
+
+    @classmethod
+    def from_file(cls, config_file: Path) -> "KernelConfigFile | None":
+        """Read and validate a config file. Returns None if the file is missing or malformed."""
+        config_data = read_config_file(config_file)
+        if config_data is None:
+            return None
+        return cls.from_dict(config_file, config_data)
+
+    def lookup_exact(self, key: AutotuneKey) -> dict[str, Any] | None:
+        if self.autotune_entries is None:
+            return None
+        return self.autotune_entries.get(AutotuneKey.key_hash(key.autotune_key))
+
+    def lookup_fuzzy(self, key: AutotuneKey) -> dict[str, Any] | None:
+        if self.autotune_entries is None:
+            return None
+        for entry in self.autotune_entries.values():
+            if key.fuzzy_matches(entry.get("autotune_key")):
+                return entry
+        return None
+
+
+@cache
+def load_config_file(config_file: Path) -> dict[str, Any] | None:
+    try:
+        with open(config_file) as f:
+            return json.load(f)
+    except Exception as e:
+        logger.warning("Error reading config file %s: %s", config_file, e)
+        return None
+
+
+def read_config_file(config_file: Path) -> dict[str, Any] | None:
+    """Read a config file, bypassing the in-process cache in ALWAYS mode."""
+    if FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
+        return load_config_file.__wrapped__(config_file)
+    return load_config_file(config_file)
+
+
+def load_cached_config(kernel_name: str, autotune_key: AutotuneKey | None = None) -> dict[str, Any] | None:
+    """
+    Load cached best config for a kernel from FLA configs directory.
+
+    This function loads the cached best configuration for a given kernel name
+    from get_fla_config_dir()/{kernel_name}.json.
+
+    Cache files may contain multiple autotune entries keyed by Triton's
+    runtime tuning key plus a top-level default config.
+
+    If the config file is not found or cannot be loaded, a warning is printed
+    and None is returned, allowing fallback to Triton's autotune.
+
+    The lookup mode is controlled by the FLA_CACHE_MODE environment variable (see FlaCacheMode).
+
+    Args:
+        kernel_name: Name of the kernel (e.g., "causal_conv1d_fwd_kernel")
+        autotune_key: Triton autotune key for the current invocation
+
+    Returns:
+        Best config dictionary or None if not found or disabled
+    """
+    if FLA_CACHE_MODE is FlaCacheMode.DISABLED:
+        return None
+
+    config_dir = get_fla_config_dir()
+    config_file = config_dir / f"{kernel_name}.json"
+
+    if not config_file.exists():
+        return None
+
+    config_data = read_config_file(config_file)
+    if config_data is None:
+        return None
+    config = KernelConfigFile.from_dict(config_file, config_data)
+    if config is None:
+        return None
+
+    if FLA_CACHE_MODE is FlaCacheMode.DEFAULT or FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
+        return config.default_config
+
+    # STRICT mode: exact match only, no fuzzy fallback
+    if FLA_CACHE_MODE is FlaCacheMode.STRICT:
+        if autotune_key is not None:
+            entry = config.lookup_exact(autotune_key)
+            if entry is not None:
+                return entry["config"]
+        return None
+
+    # FULL and FUZZY modes: try exact key match first, then fuzzy match
+    if autotune_key is not None:
+        entry = config.lookup_exact(autotune_key) or config.lookup_fuzzy(autotune_key)
+        if entry is not None:
+            return entry["config"]
+
+    if FLA_CACHE_MODE is FlaCacheMode.FUZZY:
+        return None
+
+    # FULL mode: fall back to default_config, then legacy raw config (no autotune_entries)
+    if config.default_config is not None:
+        return config.default_config
+    if config.autotune_entries is not None:
+        return None
+    return config_data
+
+
+class CachedAutotuner(Autotuner):
+    """
+    A modified autotuner that loads best config from FLA's config directory.
+
+    This class extends Triton's Autotuner but overrides the run method to
+    try loading cached configuration first before falling back to autotune.
+    """
+
+    def __init__(self, fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs):
+        super().__init__(fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs)
+        self.kernel_name = fn.fn.__name__ if hasattr(fn, 'fn') else fn.__name__
+
+        # None-safe pre/post hooks: Triton's defaults crash when a restore_value / reset_to_zero arg
+        # is None (idiomatic for optional pointers gated by a tl.constexpr flag).
+        # Fixed upstream in triton-lang/triton#10295 — remove this override once FLA's minimum Triton version has it.
+        if not self.user_defined_pre_hook and (self.reset_to_zero or self.restore_value):
+            def _pre_hook(kw, reset_only=False):
+                for n in self.reset_to_zero:
+                    if kw[n] is not None:
+                        kw[n].zero_()
+                if not reset_only:
+                    self.restore_copies = {n: kw[n].clone() for n in self.restore_value if kw[n] is not None}
+            self.pre_hook = _pre_hook
+        if not self.user_defined_post_hook and self.restore_value:
+            def _post_hook(kw, exception):
+                for n, copy in self.restore_copies.items():
+                    kw[n].copy_(copy)
+                self.restore_copies = {}
+            self.post_hook = _post_hook
+
+    def should_check_fla_cache(self, key: AutotuneKey) -> bool:
+        if FLA_CACHE_MODE is FlaCacheMode.DISABLED:
+            return False
+        if FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
+            return True
+        return key.autotune_key not in self.cache
+
+    def run(self, *args, **kwargs):
+        key = AutotuneKey.build(self.arg_names, self.keys, args, kwargs)
+        if self.should_check_fla_cache(key):
+            self.maybe_load_cached_config(key)
+        return super().run(*args, **kwargs)
+
+    def maybe_load_cached_config(self, key: AutotuneKey):
+        best_config = load_cached_config(self.kernel_name, key)
+
+        if best_config is not None:
+            kw = best_config["kwargs"]
+            num_warps = best_config["num_warps"]
+            num_stages = best_config["num_stages"]
+
+            extra = {
+                "num_ctas": best_config["num_ctas"],
+                "maxnreg": best_config.get("maxnreg"),
+                "pre_hook": None,
+                "ir_override": best_config.get("ir_override"),
+            } if TRITON_ABOVE_3_5_1 else {}
+            cfg = triton.Config(kw, num_warps=num_warps, num_stages=num_stages, **extra)
+
+            self.cache[key.autotune_key] = cfg
+        else:
+            logger.debug(
+                "No cached config found for kernel %s and key %s; falling back to Triton autotune",
+                self.kernel_name,
+                list(key.autotune_key),
+            )
+
+
+def fla_cache_autotune(configs, key=None, prune_configs_by=None, reset_to_zero=None, restore_value=None,
+                       pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False,
+                       do_bench=None, cache_results=False):
+    """
+    Decorator for auto-tuning a :code:`triton.jit`'d function with FLA config support.
+
+    Extends Triton's autotune to load best configurations from FLA's config directory
+    (default: fla/configs/{GPU}/, or FLA_CONFIG_DIR/ when overridden), keyed by kernel
+    name from {kernel_name}.json. Lookup behaviour is controlled by FLA_CACHE_MODE.
+    Falls back to normal Triton autotuning when no cached config is found.
+    """
+    # key can be None when we want to use cache only (no fallback autotune)
+    if key is None:
+        key = []
+
+    def decorator(fn):
+        kwargs = {}
+        if TRITON_ABOVE_3_4_0:
+            kwargs = {"cache_results": cache_results}
+
+        return CachedAutotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value,
+                               pre_hook=pre_hook, post_hook=post_hook,
+                               prune_configs_by=prune_configs_by, warmup=warmup, rep=rep,
+                               use_cuda_graph=use_cuda_graph, do_bench=do_bench,
+                               **kwargs,
+                               )
+
+    return decorator
+
+
+def configure_fla_cache_autotune():
+    triton.autotune = fla_cache_autotune
+    logger.info(
+        "configure_fla_cache_autotune() is enabling FLA fla_cache_autotune; "
+        "triton.autotune will be replaced with fla_cache_autotune."
+    )
+
+
+def restore_autotune_backend():
+    from triton.runtime.autotuner import autotune as original_autotune
+    triton.autotune = original_autotune
+    logger.info(
+        "restore_autotune_backend() is restoring Triton's original autotune; "
+        "triton.autotune will be replaced with triton.runtime.autotuner.autotune."
+    )

ex_engine/fla_kernels/utils/op.py

@@ -0,0 +1,101 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import os
+
+import triton
+import triton.language as tl
+import triton.language.extra.libdevice as tldevice
+
+from fla.utils import IS_GATHER_SUPPORTED, IS_NVIDIA_BLACKWELL
+
+if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
+    @triton.jit
+    def exp(x): return tldevice.fast_expf(x.to(tl.float32))
+    @triton.jit
+    def exp2(x): return tldevice.exp2(x.to(tl.float32))
+    @triton.jit
+    def log(x): return tldevice.fast_logf(x.to(tl.float32))
+    @triton.jit
+    def log2(x): return tldevice.fast_log2f(x.to(tl.float32))
+    @triton.jit
+    def tanh(x): return tldevice.fast_tanhf(x.to(tl.float32))
+else:
+    @triton.jit
+    def exp(x): return tl.exp(x.to(tl.float32))
+    @triton.jit
+    def exp2(x): return tl.math.exp2(x.to(tl.float32))
+    @triton.jit
+    def log(x): return tl.log(x.to(tl.float32))
+    @triton.jit
+    def log2(x): return tl.log2(x.to(tl.float32))
+    @triton.jit
+    def tanh(x): return tldevice.tanh(x.to(tl.float32))
+
+
+if IS_NVIDIA_BLACKWELL:
+    """
+    Compute tl.dot with Blackwell workaround.
+
+    On SM100 datacenter and SM120 consumer Blackwell GPUs, wraps the result in
+    inline assembly to prevent the TritonGPUHoistTMEMAlloc pass from incorrectly
+    fusing add and dot operations.
+    See: https://github.com/fla-org/flash-linear-attention/issues/638
+
+    TODO: Remove this workaround once the Triton compiler bug is fixed.
+    Track upstream issue at: https://github.com/triton-lang/triton/issues/8695
+    """
+    @triton.jit
+    def safe_dot(a, b, allow_tf32: tl.constexpr = None):
+        return tl.inline_asm_elementwise(
+            asm="mov.f32 $0, $1;",
+            constraints="=r,r",
+            args=[tl.dot(a, b, allow_tf32=allow_tf32)],
+            dtype=tl.float32,
+            is_pure=True,
+            pack=1,
+        )
+else:
+    @triton.jit
+    def safe_dot(a, b, allow_tf32: tl.constexpr = None):
+        return tl.dot(a, b, allow_tf32=allow_tf32)
+
+
+if not IS_GATHER_SUPPORTED:
+    @triton.jit
+    def gather(src, index, axis, _builder=None):
+        """
+        Gather operation that works when tl.gather is not supported.
+        This is a fallback implementation that returns None.
+        Just to make triton compiler happy.
+        """
+        return None
+else:
+    gather = tl.gather
+
+
+if hasattr(triton.language, '_experimental_make_tensor_descriptor'):
+    # For Triton 3.3.x
+    make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor
+elif hasattr(triton.language, 'make_tensor_descriptor'):
+    # For Triton 3.4.x and later
+    make_tensor_descriptor = triton.language.make_tensor_descriptor
+else:
+    """
+    Fallback implementation when TMA is not supported.
+    Returns None to indicate TMA descriptors are unavailable.
+    Just make triton compiler happy.
+    """
+    @triton.jit
+    def make_tensor_descriptor(
+        base,
+        shape,
+        strides,
+        block_shape,
+        _builder=None,
+    ):
+        return None

ex_engine/kernels/kernels.h

@@ -0,0 +1,11 @@
+/* Auto-generated aggregation header for xllm::kernel namespace.
+ * Equivalent to CMake cc_library(NAME kernels HDRS param.h ops_api.h).
+ * 
+ * AST Layer 3: kernel dispatch interface
+ * Called by: xllm_layers/ (Layer 2)
+ * Calls:     xllm_kernels/ilu/ (Layer 4)
+ */
+#pragma once
+
+#include "param.h"
+#include "ops_api.h"

ex_engine/kernels/ops_api.h

@@ -0,0 +1,177 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include "param.h"
+
+namespace xllm::kernel {
+
+static const std::string kActModeSilu = "silu";
+static const std::string kActModeGelu = "gelu";
+static const std::string kActModeQuickGelu = "quick_gelu";
+static const std::string kActModeSwish = "swish";
+
+void apply_rotary(RotaryParams& params);
+
+void active(ActivationParams& params);
+
+void reshape_paged_cache(ReshapePagedCacheParams& params);
+
+void reshape_from_cache(ReshapeFromCacheParams& params);
+
+// Quantize and store KV cache to paged cache (INT8 quantization)
+// Only supported on MLU backend
+void quant_to_paged_cache(ReshapePagedCacheParams& params);
+
+// Dequantize KV cache from paged cache (INT8 to FP16/BF16)
+// Only supported on MLU backend
+void dequant_from_paged_cache(ReshapeFromCacheParams& params);
+
+void fused_layernorm(FusedLayerNormParams& params);
+
+torch::Tensor matmul(MatmulParams& params);
+
+torch::Tensor group_gemm(GroupGemmParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    MoeFusedTopkParams& params);
+
+std::vector<torch::Tensor> moe_gen_idx(MoeGenIdxParams& params);
+
+torch::Tensor moe_expand_input(MoeExpandInputParams& params);
+
+torch::Tensor moe_combine_result(MoeCombineResultParams& params);
+
+torch::Tensor moe_all2all_gen_send_layout(
+    MoeAll2AllGenSendLayoutParams& params);
+
+std::vector<torch::Tensor> moe_all2all_gen_gather_index(
+    MoeAll2AllGenGatherIndexParams& params);
+
+std::vector<torch::Tensor> moe_all2all_create(MoeAll2AllCreateParams& params);
+
+void moe_all2all_init(MoeAll2AllInitParams& params);
+
+void moe_all2all_dispatch(MoeAll2AllDispatchParams& params);
+
+void moe_all2all_combine(MoeAll2AllCombineParams& params);
+
+void moe_all2all_destroy(MoeAll2AllDestroyParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor> scaled_quantize(
+    ScaledQuantizeParams& params);
+
+torch::Tensor scaled_matmul(ScaledMatmulParams& params);
+
+torch::Tensor apply_top_k_top_p(TopKPParams& params);
+
+torch::Tensor random_sample(RandomSampleParams& params);
+
+torch::Tensor rejection_sample(RejectionSampleParams& params);
+
+void masked_indexer_select_paged_kv(MaskedIndexerSelectPagedKVParams& params);
+
+void gather_split(GatherSplitParams& params);
+
+void fused_mla_q(FusedMlaQParams& params);
+
+void fused_mla_kv(FusedMlaKVParams& params);
+
+void fused_indexer_q(FusedIndexerQParams& params);
+
+void fused_indexer_k(FusedIndexerKParams& params);
+
+// L2 normalization along the last dimension
+torch::Tensor l2_norm(torch::Tensor& x, double eps = 1e-6);
+
+// TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx + moe_expand_input
+// (and token_count/cusum outputs) on other backends.
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+moe_init_routing_v2(MoeInitRoutingV2Params& params);
+
+// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format
+// Returns: (quantized_output, scale)
+std::tuple<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
+    Fp8ScaledQuantizeParams& params);
+
+// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels
+// Performs: c = (a @ b.T) with scales applied
+torch::Tensor fp8_scaled_matmul(Fp8ScaledMatmulParams& params);
+
+// Static scaled FP8 quantization helper
+// Quantizes input tensor to FP8 using a pre-computed scale factor
+void static_scaled_fp8_quant(StaticScaledFp8QuantParams& params);
+
+// Fused RMSNorm + Static FP8 Quantization
+// These fused operations combine RMSNorm and FP8 quantization to reduce memory
+// bandwidth by avoiding the intermediate write-back to global memory.
+
+// Fused RMSNorm + Static FP8 Quantization
+// Returns: FP8 quantized output tensor
+torch::Tensor rms_norm_static_fp8_quant(RmsNormStaticFp8QuantParams& params);
+
+// Fused Add + RMSNorm + Static FP8 Quantization (with residual)
+// Returns: tuple of (FP8 quantized output, updated residual)
+std::tuple<torch::Tensor, torch::Tensor> fused_add_rms_norm_static_fp8_quant(
+    FusedAddRmsNormStaticFp8QuantParams& params);
+
+std::pair<torch::Tensor, torch::Tensor> fused_gdn_gating(
+    FusedGdnGatingParams& params);
+
+std::pair<torch::Tensor, torch::Tensor> fused_recurrent_gated_delta_rule(
+    FusedRecurrentGatedDeltaRuleParams& params);
+
+torch::Tensor causal_conv1d_update(CausalConv1dUpdateParams& params);
+
+torch::Tensor gated_layer_norm(GatedLayerNormParams& params);
+
+std::pair<torch::Tensor, torch::Tensor> partial_rotary_embedding(
+    PartialRotaryEmbeddingParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+fused_qkvzba_split_reshape_cat(FusedQkvzbaSplitReshapeParams& params);
+
+void gemma_rms_norm(GemmaRMSNormParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+split_qkv_rmsnorm_mrope(SplitQkvRmsnormMropeParams& params);
+
+bool has_split_qkv_rmsnorm_mrope_specialization(int64_t num_q_heads,
+                                                int64_t num_kv_heads,
+                                                int64_t head_size);
+
+torch::Tensor build_split_qkv_rmsnorm_mrope_gather_pattern(
+    int64_t rope_dim,
+    const std::vector<int64_t>& mrope_section,
+    bool is_interleaved,
+    const torch::Device& device);
+
+std::pair<torch::Tensor, torch::Tensor> chunk_gated_delta_rule(
+    ChunkGatedDeltaRuleParams& params);
+
+torch::Tensor recurrent_gated_delta_rule(
+    const torch::Tensor& query,
+    const torch::Tensor& key,
+    const torch::Tensor& value,
+    torch::Tensor& state,
+    const std::optional<torch::Tensor>& beta,
+    const std::optional<double> scale,
+    const std::optional<torch::Tensor>& actual_seq_lengths,
+    const std::optional<torch::Tensor>& ssm_state_indices,
+    const std::optional<torch::Tensor>& num_accepted_tokens,
+    const std::optional<torch::Tensor>& g,
+    const std::optional<torch::Tensor>& gk);
+}  // namespace xllm::kernel