merge updates for Mi300x ops (#267)

* add initial r1 ops for testing

* add deepseek_r1_sigmoid_top_k_f32

* add glu_expert_bf16xf8_block_scal

* fix glu_expert_bf16xf8_block_scal for cat([gate, up], dim=0)

* update cached bf16xf8_block_scal

* add gemm_nt_bf16xfp8_block_scal

* add modules

* enhance glu perf for bs=32

* fuse routed_scaled in topk_gating

* using rotary_lookup_bf16 instead of rotary_emb_bf16

* moving head/tail ops to C++

* add partial absorb fusion

* using gate_gemm_out_bf16

* add gemm_gate_up_silu_bf16xf8_s_16x16

* use torch::matmul for some gemm

* handle scaling at non-ending dim

* add glu_expert_bf16xf8_block_scal_16x16 back

* fine tune bf16 deviation

* add test_allreduce_bf16

* update deepseek_sigmoid_top_8_static_v2 with scaling 2.5

* add glu_expert_bf16xf8_block_scal_16x16_fnuz

* restore previous interface to avoid compatiblity break

* add multi_head_latent_rope_bf16

* add system.from_url()

* force deepseek_sigmoid_top_8_static_v2's dtype compatible with sglang

* isolate NCCL-dependent APIs with macros

Author: ghostplant

MANIFEST.in

@@ -0,0 +1 @@
+include tutel/ops/*

setup.py

@@ -113,6 +113,7 @@ def install(use_cuda, use_nccl):
         install_requires=[
             "numpy",
         ],
+        include_package_data=True,
         zip_safe=False,
         extras_require={
             'test': [

tutel/custom/antares_ops.h

@@ -181,7 +181,7 @@ at::Tensor call(const void *key, const std::vector<at::Tensor> &ts, const std::v
       static std::unordered_map<std::string, decltype(torch::kInt8)> key_to_dtype = {
         {"int8", torch::kInt8}, {"int16", torch::kInt16}, {"int32", torch::kInt32}, {"int64", torch::kInt64},
         {"bfloat8", at::kFloat8_e5m2}, {"float8", at::kFloat8_e4m3fn}, {"bfloat16", torch::kBFloat16}, {"float16", torch::kFloat16}, {"float32", torch::kFloat32}, {"float64", torch::kFloat64},
-        {"bfloat2x16", at::kComplexHalf}, {"float2x16", at::kComplexHalf}, {"float2x32", at::kComplexFloat},
+        {"bfloat2x16", torch::kInt32}, {"float2x16", torch::kInt32}, {"float2x32", torch::kInt64},
       };
 
       auto dtype_it = key_to_dtype.find(o_type[1]);

tutel/custom/custom_kernel.cpp

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import torch
+import time
+import argparse
+
+from tutel import system, net
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
+parser.add_argument('--count', type=int, default=229376)
+parser.add_argument('--loop', type=int, default=50)
+parser.add_argument('--warmup', type=int, default=5, help='Number of warmup iterations')
+args = parser.parse_args()
+
+parallel_env = system.init_data_model_parallel(backend='nccl' if args.device == 'cuda' else 'gloo')
+local_device = parallel_env.local_device
+
+x = torch.randn([args.count], device=local_device, dtype=torch.float32)
+
+if args.device == 'cuda':
+  wait = lambda: torch.cuda.synchronize() or time.perf_counter()
+else:
+  wait = lambda: time.perf_counter()
+
+# Warmup phase (excluded from any measurement)
+with torch.no_grad():
+  for _ in range(args.warmup + args.loop):
+    torch.ops.tutel_ops.test_allreduce_bf16(args.count)
+

tutel/examples/nccl_allreduce_perf.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import os
+import tutel_custom_kernel
+
+if 'OP_LOADER' not in os.environ:
+    os.environ['OP_LOADER'] = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.')