[kunlunxin] fix outer,mv,scatter,batch_norm; add reshape_and_cache (#723)

* [kunlunxin] fix outer,mv,scatter,batch_norm; add reshape_and_cache

* [kunlunxin] add concat_and_cache_mla

Author: duanyaqi

src/flag_gems/runtime/backend/_kunlunxin/fused/concat_and_cache_mla.py

@@ -40,53 +40,51 @@ def concat_and_cache_mla_kernel(
     slot_idx = tl.load(slot_mapping_ptr + token_idx)
 
     # Skip padded tokens
-    if slot_idx < 0:
-        return
-
-    # Calculate cache position
-    block_id = slot_idx // block_size
-    block_offset = slot_idx % block_size
-    cache_base = block_id * block_stride + block_offset * entry_stride
-
-    # Preload scale if needed
-    if kv_dtype != FP8_KV_CACHE_DATA_TYPE_AUTO:
-        scale_val = tl.load(scale_ptr)
-
-    # Process kv_c section
-    for i in range(0, kv_lora_rank, BLOCK_SIZE):
-        idx = i + tl.arange(0, BLOCK_SIZE)
-        mask = idx < kv_lora_rank
-
-        src_ptr = kv_c_ptr + token_idx * kv_c_stride + idx
-        dst_ptr = kv_cache_ptr + cache_base + idx
-
-        val = tl.load(src_ptr, mask=mask, other=0)
-
-        if kv_dtype != FP8_KV_CACHE_DATA_TYPE_AUTO:
-            if kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E4M3:
-                val = (val / scale_val).to(tl.float8e4nv)
-            elif kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E5M2:
-                val = (val / scale_val).to(tl.float8e5)
-            val = val.to(tl.uint8, bitcast=True)
-        tl.store(dst_ptr, val, mask=mask)
-
-    # Process k_pe section
-    for j in range(0, pe_dim, BLOCK_SIZE):
-        idx = j + tl.arange(0, BLOCK_SIZE)
-        mask = idx < pe_dim
-
-        src_ptr = k_pe_ptr + token_idx * k_pe_stride + idx
-        dst_ptr = kv_cache_ptr + cache_base + kv_lora_rank + idx
-
-        val = tl.load(src_ptr, mask=mask, other=0)
+    if slot_idx >= 0:
+        # Calculate cache position
+        block_id = slot_idx // block_size
+        block_offset = slot_idx % block_size
+        cache_base = block_id * block_stride + block_offset * entry_stride
 
+        # Preload scale if needed
         if kv_dtype != FP8_KV_CACHE_DATA_TYPE_AUTO:
-            if kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E4M3:
-                val = (val / scale_val).to(tl.float8e4nv)
-            elif kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E5M2:
-                val = (val / scale_val).to(tl.float8e5)
-            val = val.to(tl.uint8, bitcast=True)
-        tl.store(dst_ptr, val, mask=mask)
+            scale_val = tl.load(scale_ptr)
+
+        # Process kv_c section
+        for i in range(0, kv_lora_rank, BLOCK_SIZE):
+            idx = i + tl.arange(0, BLOCK_SIZE)
+            mask = idx < kv_lora_rank
+
+            src_ptr = kv_c_ptr + token_idx * kv_c_stride + idx
+            dst_ptr = kv_cache_ptr + cache_base + idx
+
+            val = tl.load(src_ptr, mask=mask, other=0)
+
+            if kv_dtype != FP8_KV_CACHE_DATA_TYPE_AUTO:
+                if kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E4M3:
+                    val = (val / scale_val).to(tl.float8e4nv)
+                elif kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E5M2:
+                    val = (val / scale_val).to(tl.float8e5)
+                val = val.to(tl.uint8, bitcast=True)
+            tl.store(dst_ptr, val, mask=mask)
+
+        # Process k_pe section
+        for j in range(0, pe_dim, BLOCK_SIZE):
+            idx = j + tl.arange(0, BLOCK_SIZE)
+            mask = idx < pe_dim
+
+            src_ptr = k_pe_ptr + token_idx * k_pe_stride + idx
+            dst_ptr = kv_cache_ptr + cache_base + kv_lora_rank + idx
+
+            val = tl.load(src_ptr, mask=mask, other=0)
+
+            if kv_dtype != FP8_KV_CACHE_DATA_TYPE_AUTO:
+                if kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E4M3:
+                    val = (val / scale_val).to(tl.float8e4nv)
+                elif kv_dtype == FP8_KV_CACHE_DATA_TYPE_FP8E5M2:
+                    val = (val / scale_val).to(tl.float8e5)
+                val = val.to(tl.uint8, bitcast=True)
+            tl.store(dst_ptr, val, mask=mask)
 
 
 class ConcatAndCacheMla(torch.autograd.Function):

src/flag_gems/runtime/backend/_kunlunxin/fused/outer.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from ..ops import mul, mv
+from ..ops import mul, mv_cluster
 
 logger = logging.getLogger(__name__)
 
@@ -27,8 +27,8 @@ def backward(ctx, out_grad):
 
         inp, weight = ctx.saved_tensors
 
-        inp_grad = mv(out_grad, weight)
-        weight_grad = mv(out_grad.t().contiguous(), inp)
+        inp_grad = mv_cluster(out_grad, weight)
+        weight_grad = mv_cluster(out_grad.t().contiguous(), inp)
 
         return inp_grad, weight_grad
 

src/flag_gems/runtime/backend/_kunlunxin/fused/reshape_and_cache.py

@@ -29,41 +29,39 @@ def reshape_and_cache_kernel(
 ):
     token_idx = tl.program_id(0)
     slot_idx = tl.load(slot_mapping + token_idx)
-    if slot_idx < 0:
-        return
+    if slot_idx >= 0:
+        block_idx = slot_idx // block_size
+        block_offset = slot_idx % block_size
+        i = tl.arange(0, triton.next_power_of_2(n))
+        mask = i < n
 
-    block_idx = slot_idx // block_size
-    block_offset = slot_idx % block_size
-    i = tl.arange(0, triton.next_power_of_2(n))
-    mask = i < n
+        src_key_idx = token_idx * key_stride + i
+        src_value_idx = token_idx * value_stride + i
+        head_idx = i // head_size
+        head_offset = i % head_size
+        x_idx = head_offset // x
+        x_offset = head_offset % x
 
-    src_key_idx = token_idx * key_stride + i
-    src_value_idx = token_idx * value_stride + i
-    head_idx = i // head_size
-    head_offset = i % head_size
-    x_idx = head_offset // x
-    x_offset = head_offset % x
-
-    tgt_key_idx = (
-        block_idx * num_heads * (head_size // x) * block_size * x
-        + head_idx * (head_size // x) * block_size * x
-        + x_idx * block_size * x
-        + block_offset * x
-        + x_offset
-    )
-    tgt_value_idx = (
-        block_idx * num_heads * head_size * block_size
-        + head_idx * head_size * block_size
-        + head_offset * block_size
-        + block_offset
-    )
+        tgt_key_idx = (
+            block_idx * num_heads * (head_size // x) * block_size * x
+            + head_idx * (head_size // x) * block_size * x
+            + x_idx * block_size * x
+            + block_offset * x
+            + x_offset
+        )
+        tgt_value_idx = (
+            block_idx * num_heads * head_size * block_size
+            + head_idx * head_size * block_size
+            + head_offset * block_size
+            + block_offset
+        )
 
-    tgt_key = tl.load(key + src_key_idx, mask=mask)
-    tgt_value = tl.load(value + src_value_idx, mask=mask)
+        tgt_key = tl.load(key + src_key_idx, mask=mask)
+        tgt_value = tl.load(value + src_value_idx, mask=mask)
 
-    # TODO: support fp8 dtype
-    tl.store(key_cache + tgt_key_idx, tgt_key, mask=mask)
-    tl.store(value_cache + tgt_value_idx, tgt_value, mask=mask)
+        # TODO: support fp8 dtype
+        tl.store(key_cache + tgt_key_idx, tgt_key, mask=mask)
+        tl.store(value_cache + tgt_value_idx, tgt_value, mask=mask)
 
 
 def reshape_and_cache(

src/flag_gems/runtime/backend/_kunlunxin/ops/__init__.py

@@ -104,7 +104,7 @@
 from .mse_loss import mse_loss
 from .mul import mul, mul_
 from .multinomial import multinomial
-from .mv import mv
+from .mv import mv, mv_cluster
 from .nan_to_num import nan_to_num
 from .ne import ne, ne_scalar
 from .neg import neg, neg_
@@ -311,6 +311,7 @@
     "normal_tensor_tensor",
     "uniform_",
     "mv",
+    "mv_cluster",
     "nan_to_num",
     "ne",
     "ne_scalar",

src/flag_gems/runtime/backend/_kunlunxin/ops/batch_norm.py

@@ -37,11 +37,11 @@ def make_3d_for_bn(input: Tensor) -> Tensor:
 
 
 @libentry()
-@triton.autotune(
-    configs=runtime.get_tuned_config("batch_norm"),
-    key=["batch_dim", "spatial_dim"],
-    restore_value=["running_mean_pointer", "running_var_pointer"],
-)
+# @triton.autotune(
+#     configs=runtime.get_tuned_config("batch_norm"),
+#     key=["batch_dim", "spatial_dim"],
+#     restore_value=["running_mean_pointer", "running_var_pointer"],
+# )
 @triton.heuristics(runtime.get_heuristic_config("batch_norm"))
 @triton.jit
 def batch_norm_forward_kernel(
@@ -186,10 +186,10 @@ def batch_norm_heur_block_n(args):
 
 
 @libentry()
-@triton.autotune(
-    configs=runtime.get_tuned_config("batch_norm"),
-    key=["batch_dim", "spatial_dim"],
-)
+# @triton.autotune(
+#     configs=runtime.get_tuned_config("batch_norm"),
+#     key=["batch_dim", "spatial_dim"],
+# )
 @triton.heuristics(
     values={
         "BLOCK_M": batch_norm_heur_block_m,

src/flag_gems/runtime/backend/_kunlunxin/ops/mv.py

@@ -98,3 +98,25 @@ def mv(inp, vec):
             out = out.squeeze()
             del os.environ["XMLIR_MATMUL_FAST_MODE"]
     return out
+
+
+def mv_cluster(inp, vec):
+    logger.debug("GEMS MV")
+    assert inp.shape[1] == vec.shape[0], "incompatible dimensions"
+    N, M = inp.shape
+    out = torch.empty((N,), device=inp.device, dtype=inp.dtype)
+    grid = lambda META: (triton.cdiv(N, META["BLOCK_N"]),)
+    with torch_device_fn.device(inp.device):
+        mv_kernel[grid](
+            inp,
+            vec,
+            out,
+            N,
+            M,
+            inp.stride(0),
+            inp.stride(1),
+            vec.stride(0),
+            out.stride(0),
+            buffer_size_limit=256,
+        )
+    return out

src/flag_gems/runtime/backend/_kunlunxin/ops/scatter.py

@@ -6,7 +6,7 @@
 import torch
 
 from flag_gems.utils.code_cache import code_cache_dir
-from flag_gems.utils.code_utils import IndentedBuffer
+from flag_gems.utils.code_utils import IndentedBuffer, write_atomic
 from flag_gems.utils.shape_utils import (
     MemOverlap,
     has_internal_overlapping,
@@ -23,6 +23,7 @@ def generate_imports(code: IndentedBuffer) -> IndentedBuffer:
     code.newline()
     code.writeline("from flag_gems.utils import libentry")
     code.writeline("from flag_gems import runtime")
+    code.writeline("import flag_gems")
     # code.writeline("from flag_gems.utils import triton_lang_extension as tle")
     code.newline()
     code.newline()
@@ -159,17 +160,13 @@ def generate_scatter_kernel(
             code.writeline("if IS_ADD: ")
             with code.indent():
                 code.writeline(
-                    "cur_inp = tl.load(inp + inp_offsets, mask=mask, other=0)"
+                    "tl.atomic_add(out + inp_offsets, cur_src, mask=mask, sem='relaxed')"
                 )
-                code.writeline("res = cur_inp + cur_src")
-                code.writeline("tl.store(out + inp_offsets, res, mask=mask)")
             code.writeline("elif IS_MUL: ")
             with code.indent():
                 code.writeline(
-                    "cur_inp = tl.load(inp + inp_offsets, mask=mask, other=0)"
+                    "tl.atomic_mul(out + inp_offsets, cur_src, mask=mask, sem='relaxed')"
                 )
-                code.writeline("res = cur_inp * cur_src")
-                code.writeline("tl.store(out + inp_offsets, res, mask=mask)")
 
             code.writeline("else: ")
             with code.indent():
@@ -252,6 +249,7 @@ def generate_destination_passing_wrapper(
                 code.writeline("IS_ADD,")
                 code.writeline("IS_MUL,")
                 code.writeline("INT32_OFFSET=int32_offset,")
+                # code.writeline("buffer_size_limit=512,")
                 # code.writeline("isCloseUnrollControl=True,")
 
         code.writeline(")")
@@ -294,15 +292,14 @@ def __call__(self, *args, **kwargs):
                 code,
             )
 
-            file_name = f"scatter_rank_{key}_pid_{self.pid}.py"
-
-            with open(code_cache_dir() / file_name, "wt", encoding="utf-8") as f:
-                f.write(code.getvalue())
+            file_name = f"scatter_rank_{key}.py"
+            file_path = code_cache_dir() / file_name
+            write_atomic(file_path, code.getvalue())
 
             # load
             spec = importlib.util.spec_from_file_location(
-                f"_gen_module_rank_{key}_pid_{self.pid}",
-                f.name,
+                f"_gen_module_rank_{key}",
+                file_path,
             )
 
             m = importlib.util.module_from_spec(spec)
@@ -319,7 +316,6 @@ def arg_key(self, *args):
 
 
 _scatter_func = ScatterFunction()
-_scatter_inplace_func = ScatterFunction()
 
 
 def scatter(inp, dim, index, src, reduce=None):
@@ -378,7 +374,7 @@ def scatter_(inp, dim, index, src, reduce=None):
 
     int32_size_dim = lambda x: x.stride(dim) * x.size(dim) < 2**32
     use_int32_offset = all(map(int32_size_dim, (inp, index, src)))
-    _scatter_inplace_func(
+    _scatter_func(
         src_restrided,
         index,
         inp_restrided,

tests/test_attention_ops.py

@@ -127,7 +127,6 @@ def convert_fp8(
 
 @pytest.mark.skipif(flag_gems.vendor_name == "hygon", reason="RuntimeError")
 @pytest.mark.skipif(flag_gems.device == "musa", reason="RuntimeError")
-@pytest.mark.skipif(flag_gems.vendor_name == "kunlunxin", reason="RESULT TODOFIX")
 @pytest.mark.concat_and_cache_mla
 @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
 @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)

tests/test_quant.py

@@ -35,7 +35,7 @@
 CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
 
 # We assume fp8 is always enabled for testing.
-KV_CACHE_DTYPE = ["auto", "fp8"]
+KV_CACHE_DTYPE = ["auto", "fp8"] if flag_gems.vendor_name != "kunlunxin" else ["auto"]
 
 
 def _create_mla_cache(
@@ -65,7 +65,6 @@ def convert_fp8(
 
 @pytest.mark.skipif(flag_gems.vendor_name == "hygon", reason="RuntimeError")
 @pytest.mark.skipif(flag_gems.device == "musa", reason="RuntimeError")
-@pytest.mark.skipif(flag_gems.vendor_name == "kunlunxin", reason="RESULT TODOFIX")
 @pytest.mark.concat_and_cache_mla
 @pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
 @pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)